marEx 2.0.0__py3-none-any.whl

marEx/plotX/base.py

@@ -0,0 +1,387 @@
+from pathlib import Path
+import numpy as np
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
+import dask
+from PIL import Image
+import subprocess
+from matplotlib.colors import ListedColormap, BoundaryNorm
+from dataclasses import dataclass
+from typing import Optional, Dict, List, Union, Tuple
+import shutil
+import warnings
+import os
+
+@dataclass
+class PlotConfig:
+    """Configuration class for plot parameters"""
+    title: Optional[str] = None
+    var_units: str = ''
+    issym: bool = False
+    cmap: Optional[str] = None
+    cperc: List[int] = None
+    clim: Optional[Tuple[float, float]] = None
+    show_colorbar: bool = True
+    grid_lines: bool = True
+    grid_labels: bool = False
+    dimensions: Dict[str, str] = None
+    norm: Optional[object] = None
+    plot_IDs: bool = False
+    extend: str = 'both'
+    
+    def __post_init__(self):
+        if self.cperc is None:
+            self.cperc = [4, 96]
+        if self.dimensions is None:
+            self.dimensions = {'time': 'time', 'ydim': 'lat', 'xdim': 'lon'}
+        if self.plot_IDs:
+            self.show_colorbar = False
+
+class PlotterBase:
+    def __init__(self, xarray_obj):
+        self.da = xarray_obj
+        
+        # Cache common features
+        self._land = cfeature.LAND.with_scale('50m')
+        self._coastlines = cfeature.COASTLINE.with_scale('50m')
+    
+    def _setup_common_params(self, config: PlotConfig) -> Tuple:
+        """Centralise common parameter setup"""
+        self.setup_plot_params()
+        
+        if config.plot_IDs:
+            cmap, norm, var_units = self.setup_id_plot_params(config.cmap)
+            clim = None
+            extend = 'neither'
+            self.da = self.da.where(self.da > 0)  # Fill value to NaN (get rid of 0s)
+        else:
+            if config.cmap is None:
+                cmap = 'RdBu_r' if config.issym else 'viridis'
+            else:
+                cmap = config.cmap
+            norm = config.norm
+            if config.clim is None and norm is None:
+                clim = self.clim_robust(self.da.values, config.issym, config.cperc)
+            else:
+                clim = config.clim
+            var_units = config.var_units
+            extend = config.extend
+                
+        return cmap, norm, clim, var_units, extend
+
+    def _setup_axes(self, ax=None):
+        """Create or use existing axes with projection"""
+        if ax is None:
+            fig = plt.figure(figsize=(7, 5))
+            ax = plt.axes(projection=ccrs.Robinson())
+        else:
+            fig = ax.get_figure()
+        return fig, ax
+
+    def _add_map_features(self, ax, grid_lines=True, grid_labels=True):
+        """Add common map features to the plot"""
+        ax.add_feature(self._land, facecolor='darkgrey', zorder=2)
+        ax.add_feature(self._coastlines, linewidth=0.5, zorder=3)
+        if grid_lines:
+            ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=grid_labels,
+                        linewidth=1, color='gray', alpha=0.5, linestyle='--', zorder=4)
+
+    def _setup_colorbar(self, fig, im, show_colorbar, var_units, extend='both', 
+                       position=None):
+        """Set up colorbar with common parameters"""
+        if not show_colorbar:
+            return None
+            
+        if position is not None:
+            # For column plots
+            cbar_ax = fig.add_axes(position)
+            cb = fig.colorbar(im, cax=cbar_ax, extend=extend)
+        else:
+            # For single plots
+            cb = plt.colorbar(im, shrink=0.6, ax=plt.gca(), extend=extend)
+            
+        if var_units:
+            cb.ax.set_ylabel(var_units, fontsize=10)
+        cb.ax.tick_params(labelsize=10)
+        return cb
+
+    def _get_title(self, time_index, col_name, dimensions):
+        """Generate appropriate title based on dimension"""
+        if col_name == dimensions['time']:
+            return f"{self.da[col_name].isel({col_name: time_index}).time.dt.strftime('%Y-%m-%d').values}"
+        return f"{col_name}={self.da[col_name].isel({col_name: time_index}).values}"
+
+
+    def single_plot(self, config: PlotConfig, ax=None):
+        """Make a single plot with given configuration"""
+        cmap, norm, clim, var_units, extend = self._setup_common_params(config)
+        
+        fig, ax = self._setup_axes(ax)
+        
+        # Call implementation-specific plot function
+        ax, im = self.plot(ax=ax, cmap=cmap, clim=clim, norm=norm)
+        
+        if config.title:
+            ax.set_title(config.title, size=12)
+        
+        self._setup_colorbar(fig, im, config.show_colorbar, var_units, extend)
+        self._add_map_features(ax, config.grid_lines, config.grid_labels)
+        
+        return fig, ax
+
+    def multi_plot(self, config: PlotConfig, col='time', col_wrap=3):
+        """Make wrapped subplots with given configuration"""
+        npanels = self.da[col].size
+        nrows = int(np.ceil(npanels/col_wrap))
+        ncols = min(npanels, col_wrap)
+        
+        cmap, norm, clim, var_units, extend = self._setup_common_params(config)
+        
+        fig = plt.figure(figsize=(6*ncols, 3*nrows))
+        axes = fig.subplots(nrows, ncols, 
+                          subplot_kw={'projection': ccrs.Robinson()}).flatten()
+        
+        # Create a single plotter instance to be reused
+        base_plotter = type(self)(self.da)
+        for attr in ['fpath_tgrid', 'fpath_ckdtree']:
+            if hasattr(self, attr):
+                setattr(base_plotter, attr, getattr(self, attr))
+        
+        for i, ax in enumerate(axes):
+            if i < npanels:
+                title = self._get_title(i, col, config.dimensions)
+                
+                # Create new config for individual panel
+                panel_config = PlotConfig(
+                    title=title, cmap=cmap, clim=clim,
+                    show_colorbar=False, grid_labels=False,
+                    norm=norm, plot_IDs=False, extend=extend
+                )
+                
+                # Update data in base plotter instead of creating new instance
+                base_plotter.da = self.da.isel({col: i})
+                
+                # Plot individual panel using the same plotter instance
+                base_plotter.single_plot(panel_config, ax=ax)
+            else:
+                fig.delaxes(ax)
+        
+        # Add single colorbar for all panels
+        if config.show_colorbar:
+            fig.subplots_adjust(right=0.9)
+            if norm is None and clim is not None:
+                # Create a proper norm from clim
+                from matplotlib.colors import Normalize
+                norm = Normalize(vmin=clim[0], vmax=clim[1])
+            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+            sm.set_array([])
+            self._setup_colorbar(fig, sm, True, var_units, extend,
+                            position=[0.92, 0.15, 0.02, 0.7])
+        
+        return fig, axes
+    
+
+    def animate(self, config: PlotConfig, plot_dir='./', file_name=None):
+        """Create an animation from time series data"""
+        
+        # Check if ffmpeg is installed
+        if shutil.which('ffmpeg') is None:
+            warnings.warn(
+                "ffmpeg executable not found in system PATH. Cannot create animation.\n"
+                "Please install ffmpeg using one of the following methods:\n"
+                "  - Linux: sudo apt install ffmpeg (Ubuntu/Debian) or sudo yum install ffmpeg (CentOS/RHEL)\n"
+                "  - Conda: conda install -c conda-forge ffmpeg\n"
+                "Alternatively, use matplotlib for animation in Jupyter notebooks."
+            )
+            return None
+        
+        plot_dir = Path(plot_dir)
+        plot_dir.mkdir(exist_ok=True)
+        temp_dir = plot_dir / "blobs_seq"
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        temp_dir.mkdir(exist_ok=True)
+        if not file_name:
+            file_name = f"movie_{self.da.name}.mp4"
+        
+        output_file = plot_dir / f'{file_name}.mp4'
+        
+        # Set up plotting parameters
+        cmap, norm, clim, var_units, extend = self._setup_common_params(config)
+        
+        plot_params = {
+            'cmap': cmap,
+            'norm': norm,
+            'clim': clim,
+            'var_units': var_units,
+            'extend': extend,
+            'show_colorbar': config.show_colorbar
+        }
+        
+        # Set up grid information if needed
+        grid_info = None
+        if hasattr(self, 'fpath_tgrid') or hasattr(self, 'fpath_ckdtree'):
+            grid_info = {
+                'type': 'unstructured',
+                'tgrid_path': getattr(self, 'fpath_tgrid', None),
+                'ckdtree_path': getattr(self, 'fpath_ckdtree', None),
+                'res': 0.3
+            }
+        
+        # Generate frames using dask for parallel processing
+        delayed_tasks = []
+        for time_ind in range(len(self.da[config.dimensions['time']])):
+            data_slice = self.da.isel({config.dimensions['time']: time_ind})
+            plot_params['time_str'] = str(
+                self.da[config.dimensions['time']]
+                .isel({config.dimensions['time']: time_ind})
+                .dt.strftime('%Y-%m-%d').values
+            )
+            delayed_tasks.append(
+                make_frame(data_slice, time_ind, temp_dir, plot_params, grid_info)
+            )
+        
+        filenames = dask.compute(*delayed_tasks)
+        filenames = sorted(filenames, key=lambda x: int(x.split('_')[-1].split('.')[0]))
+        
+        # Create movie using ffmpeg
+        subprocess.run([
+            'ffmpeg', '-y', '-threads', '0', '-framerate', '10',
+            '-i', str(temp_dir / 'time_%04d.jpg'),
+            '-c:v', 'libx264', '-preset', 'fast', '-crf', '22',
+            '-pix_fmt', 'yuv420p', '-movflags', '+faststart',
+            str(output_file)
+        ], check=True)
+        
+        return str(output_file)
+
+
+    def clim_robust(self, data, issym, percentiles=[2, 98]):
+        """Base method for computing colour limits"""
+        clim = np.nanpercentile(data, percentiles)
+        
+        if issym:
+            clim = np.abs(clim).max()
+            clim = np.array([-clim, clim])
+        elif percentiles[0] == 0:
+            clim = np.array([0, clim[1]])
+
+        return clim
+    
+    def setup_plot_params(self):
+        """Set up common plotting parameters"""
+        plt.rc('text', usetex=False)
+        plt.rc('font', family='serif')
+
+    def setup_id_plot_params(self, cmap=None):
+        """Set up parameters for plotting IDs"""
+        unique_values = np.unique(self.da.values[~np.isnan(self.da.values)])
+        unique_values = unique_values[unique_values > 0]
+        bounds = np.arange(unique_values.min(), unique_values.max() + 2) - 0.5
+        n_bins = len(bounds) - 1
+        
+        if cmap is None:
+            np.random.seed(42)
+            cmap = ListedColormap(np.random.random(size=(n_bins, 3)))
+        
+        norm = BoundaryNorm(bounds, cmap.N)
+        return cmap, norm, 'ID'
+
+
+
+@dask.delayed
+def make_frame(data_slice, time_ind, temp_dir, plot_params, grid_info=None):
+    """Create a single frame for movies - minimise memory usage with dask
+    
+    Args:
+        data_slice: The data for this specific frame
+        time_ind: Frame index
+        temp_dir: Directory for temporary files
+        plot_params: Dict containing plotting parameters
+        grid_info: Dict containing grid paths and settings for unstructured data
+    """
+    # Set up plotting parameters
+    plt.rc('text', usetex=False)
+    plt.rc('font', family='serif')
+    
+    fig = plt.figure(figsize=(7, 5))
+    ax = plt.axes(projection=ccrs.Robinson())
+    
+    data_slice_np = data_slice.values
+        
+    # Set up plot kwargs
+    plot_kwargs = {
+        'transform': ccrs.PlateCarree(),
+        'cmap': plot_params['cmap'],
+        'shading': 'auto'
+    }
+    
+    if plot_params.get('norm'):
+        plot_kwargs['norm'] = plot_params['norm']
+    elif plot_params.get('clim'):
+        plot_kwargs['vmin'] = plot_params['clim'][0]
+        plot_kwargs['vmax'] = plot_params['clim'][1]
+    
+    # Handle different grid types
+    if grid_info and grid_info.get('type') == 'unstructured':
+        from .unstructured import _load_ckdtree, _load_triangulation
+        if grid_info.get('ckdtree_path'):
+            # Use cached ckdtree data
+            ckdt_data = _load_ckdtree(grid_info['ckdtree_path'], grid_info.get('res', 0.3))
+            grid_data = data_slice_np[ckdt_data['indices']].reshape(
+                ckdt_data['lat'].size,
+                ckdt_data['lon'].size
+            )
+            grid_data = np.ma.masked_invalid(grid_data)
+            im = ax.pcolormesh(
+                ckdt_data['lon'], ckdt_data['lat'],
+                grid_data, **plot_kwargs
+            )
+        elif grid_info.get('tgrid_path'):
+            # Use triangulation
+            triang = _load_triangulation(grid_info['tgrid_path'])
+            data_masked = np.ma.masked_invalid(data_slice_np)
+            im = ax.tripcolor(triang, data_masked, **plot_kwargs)
+    else:
+        # Regular grid plotting
+        lat = data_slice.lat.values
+        lon = data_slice.lon.values
+        im = ax.pcolormesh(
+                lon, lat,
+                data_slice_np, **plot_kwargs
+            )
+    
+    time_str = plot_params.get('time_str', f'Frame {time_ind}')
+    ax.set_title(time_str, size=12)
+    
+    if plot_params.get('show_colorbar'):
+        cb = plt.colorbar(im, shrink=0.6, ax=ax, extend=plot_params.get('extend', 'both'))
+        if plot_params.get('var_units'):
+            cb.ax.set_ylabel(plot_params['var_units'], fontsize=10)
+        cb.ax.tick_params(labelsize=10)
+    
+    land = cfeature.LAND.with_scale('50m')
+    coastlines = cfeature.COASTLINE.with_scale('50m')
+    ax.add_feature(land, facecolor='darkgrey', zorder=2)
+    ax.add_feature(coastlines, linewidth=0.5, zorder=3)
+    ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
+                linewidth=1, color='gray', alpha=0.5, linestyle='--', zorder=4)
+    
+    # Save and process frame
+    filename = f'time_{time_ind:04d}.jpg'
+    temp_file = temp_dir / f'temp_{filename}'
+    fig.savefig(str(temp_file), dpi=300, bbox_inches='tight')
+    plt.close(fig)
+    
+    # Ensure dimensions are even for video encoding
+    image = Image.open(str(temp_file))
+    width, height = image.size
+    new_width = width - (width % 2)
+    new_height = height - (height % 2)
+    image = image.resize((new_width, new_height), Image.LANCZOS)
+    
+    image.save(str(temp_dir / filename))
+    temp_file.unlink()
+    
+    return filename

marEx/plotX/gridded.py

@@ -0,0 +1,55 @@
+import numpy as np
+import xarray as xr
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+
+from .base import PlotterBase
+
+class GriddedPlotter(PlotterBase):
+    def __init__(self, xarray_obj):
+        super().__init__(xarray_obj)
+    
+    def wrap_lon(self, data, dimensions):
+        """Handle periodic boundary in longitude by adding a column of data."""
+        lon = data[dimensions['xdim']]
+        
+        # Check if we're dealing with global data that needs wrapping
+        lon_spacing = np.diff(lon)[0]
+        if abs(360 - (lon.max() - lon.min())) < 2 * lon_spacing:
+            # Add a column at lon=360 that equals the data at lon=0
+            new_lon = np.append(lon, lon[0] + 360)
+            wrapped_data = xr.concat([data, data.isel({dimensions['xdim']: 0})], 
+                                   dim=dimensions['xdim'])
+            wrapped_data[dimensions['xdim']] = new_lon
+            return wrapped_data
+        return data
+
+    def plot(self, ax, cmap='viridis', clim=None, norm=None):
+        """Implement plotting for gridded (i.e. regular grid) data."""
+        dimensions = {'ydim': 'lat', 'xdim': 'lon'}
+        data = self.wrap_lon(self.da, dimensions)
+        
+         # Ensure data has only required dimensions for imshow
+        if 'time' in data.dims and len(data.time) == 1:
+            data = data.squeeze(dim='time')  # Remove time dimension if it's singular
+        
+        plot_kwargs = {
+            'transform': ccrs.PlateCarree(),
+            'cmap': cmap,
+            'shading': 'auto'
+        }
+        
+        if norm is not None:
+            plot_kwargs['norm'] = norm
+        elif clim is not None:
+            plot_kwargs['vmin'] = clim[0]
+            plot_kwargs['vmax'] = clim[1]
+        
+        lons = data[dimensions['xdim']].values
+        lats = data[dimensions['ydim']].values
+        values = data.values
+        
+        # imshow has some dimension issues with cartopy...
+        im = ax.pcolormesh(lons, lats, values, **plot_kwargs)
+        
+        return ax, im

marEx/plotX/unstructured.py

@@ -0,0 +1,147 @@
+import numpy as np
+import xarray as xr
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+from matplotlib.tri import Triangulation
+from pathlib import Path
+
+from .base import PlotterBase
+
+# Global cache for grid data
+_GRID_CACHE = {
+    'triangulation': {},  # key: grid_path, value: triangulation object
+    'ckdtree': {}         # key: (ckdtree_path, res), value: ckdtree data
+}
+
+def clear_cache():
+    """Clear the global grid cache."""
+    _GRID_CACHE['triangulation'].clear()
+    _GRID_CACHE['ckdtree'].clear()
+
+def _load_triangulation(fpath_tgrid):
+    """Load and cache triangulation data globally."""
+    fpath_tgrid = str(fpath_tgrid)  # Convert Path to string for dict key
+    if fpath_tgrid not in _GRID_CACHE['triangulation']:
+        # Only load required variables
+        grid_data = xr.open_dataset(
+            fpath_tgrid,
+            chunks={},
+            drop_variables=[v for v in xr.open_dataset(fpath_tgrid).variables 
+                          if v not in ['vertex_of_cell', 'clon', 'clat']]
+        );
+        
+        if 'vertex_of_cell' not in grid_data.variables or 'clon' not in grid_data.variables or 'clat' not in grid_data.variables:
+            raise ValueError("Triangulation Grid File must contain 'vertex_of_cell', 'clon', and 'clat' variables")
+        
+        # Extract triangulation vertices - convert to 0-based indexing
+        triangles = grid_data.vertex_of_cell.values.T - 1
+        # Create matplotlib triangulation object
+        _GRID_CACHE['triangulation'][fpath_tgrid] = Triangulation(
+            grid_data.clon.values,
+            grid_data.clat.values,
+            triangles
+        )
+        grid_data.close()
+        
+    return _GRID_CACHE['triangulation'][fpath_tgrid]
+
+def _load_ckdtree(fpath_ckdtree, res):
+    """Load and cache ckdtree data globally."""
+    cache_key = (str(fpath_ckdtree), res)  # Convert Path to string for dict key
+    
+    if cache_key not in _GRID_CACHE['ckdtree']:
+        # Format resolution string to match file naming
+        ckdtree_file = Path(fpath_ckdtree) / f"res{res:3.2f}.nc"
+        
+        if not ckdtree_file.exists():
+            raise ValueError(f"No ckdtree file found for resolution {res}")
+            
+        ds_ckdt = xr.open_dataset(ckdtree_file)
+        _GRID_CACHE['ckdtree'][cache_key] = {
+            'indices': ds_ckdt.ickdtree_c.values,
+            'lon': ds_ckdt.lon.values,
+            'lat': ds_ckdt.lat.values
+        }
+        ds_ckdt.close()
+        
+    return _GRID_CACHE['ckdtree'][cache_key]
+
+class UnstructuredPlotter(PlotterBase):
+    def __init__(self, xarray_obj):
+        super().__init__(xarray_obj)
+        
+        from . import _fpath_tgrid, _fpath_ckdtree
+        self.fpath_tgrid = _fpath_tgrid
+        self.fpath_ckdtree = _fpath_ckdtree
+        
+    def specify_grid(self, fpath_tgrid=None, fpath_ckdtree=None):
+        """Set the path to the unstructured grid files."""
+        self.fpath_tgrid = Path(fpath_tgrid) if fpath_tgrid else None
+        self.fpath_ckdtree = Path(fpath_ckdtree) if fpath_ckdtree else None
+
+    def plot(self, ax, cmap='viridis', clim=None, norm=None):
+        """Implement plotting for unstructured data."""
+        if self.fpath_ckdtree is not None:
+            # Interpolate using pre-computed KDTree indices
+            grid_lon, grid_lat, grid_data = self._interpolate_with_ckdtree(self.da.values, res=0.3)
+            
+            plot_kwargs = {
+                'transform': ccrs.PlateCarree(),
+                'cmap': cmap,
+                'shading': 'auto'
+            }
+            
+            if norm is not None:
+                plot_kwargs['norm'] = norm
+            elif clim is not None:
+                plot_kwargs['vmin'] = clim[0]
+                plot_kwargs['vmax'] = clim[1]
+            
+            # Mask NaNs
+            grid_data = np.ma.masked_invalid(grid_data)
+            im = ax.pcolormesh(grid_lon, grid_lat, grid_data, **plot_kwargs)
+            
+        else:
+            # Use triangulation from file if available
+            if self.fpath_tgrid is None:
+                raise ValueError("Either fpath_tgrid or fpath_ckdtree must be provided")
+            
+            triang = _load_triangulation(self.fpath_tgrid)
+            
+            plot_kwargs = {
+                'transform': ccrs.PlateCarree(),
+                'cmap': cmap
+            }
+            
+            if norm is not None:
+                plot_kwargs['norm'] = norm
+            elif clim is not None:
+                plot_kwargs['vmin'] = clim[0]
+                plot_kwargs['vmax'] = clim[1]
+            
+            # Mask NaNs
+            native_data = self.da.copy()
+            native_data = np.ma.masked_invalid(native_data)
+            
+            im = ax.tripcolor(triang, native_data, **plot_kwargs)
+        
+        return ax, im
+
+    def _interpolate_with_ckdtree(self, data, res):
+        """Interpolate unstructured data using pre-computed KDTree indices."""
+        if self.fpath_ckdtree is None:
+            raise ValueError("ckdtree path not specified")
+            
+        # Load or get cached ckdtree data
+        ckdt_data = _load_ckdtree(self.fpath_ckdtree, res)
+        
+        # Create meshgrid for plotting
+        grid_lon_2d, grid_lat_2d = np.meshgrid(ckdt_data['lon'], ckdt_data['lat'])
+        
+        # Use indices to create interpolated data
+        grid_data = data[ckdt_data['indices']].reshape(
+            ckdt_data['lat'].size,
+            ckdt_data['lon'].size
+        )
+        
+        return grid_lon_2d, grid_lat_2d, grid_data