Source code for utopya.eval.plots.attractor

"""This module provides plotting functions to visualize the attractive set of
a dynamical system.

.. todo::

    Migrate these to the more generic DAG-based interface.

import copy
import itertools
import logging
from typing import Callable, Dict, Sequence, Tuple, Union

import matplotlib as mpl
import numpy as np
import xarray as xr
from matplotlib.collections import PatchCollection
from matplotlib.patches import Circle, Rectangle
from scipy.signal import find_peaks

# FIXME use local import
import utopya.eval.plots._attractor as utdp

from import recursive_update
from .. import DataManager, UniverseGroup
from . import MultiversePlotCreator, PlotHelper, is_plot_func
from ._mpl import HandlerEllipse
from ._utils import calc_pxmap_rectangles

log = logging.getLogger(__name__)

# -----------------------------------------------------------------------------

[docs]@is_plot_func(creator="multiverse") def bifurcation_diagram( dm: DataManager, *, hlpr: PlotHelper, mv_data: xr.Dataset, dim: str = None, dims: Tuple[str, str] = None, analysis_steps: Sequence[Union[str, Tuple[str, str]]], custom_analysis_funcs: Dict[str, Callable] = None, analysis_kwargs: dict = None, visualization_kwargs: dict = None, to_plot: dict = None, **kwargs, ) -> None: """Plots a bifurcation diagram for one or two parameter dimensions (arguments ``dim`` or ``dims``). Args: dm (DataManager): The data manager from which to retrieve the data hlpr (PlotHelper): The PlotHelper that instantiates the figure and takes care of plot aesthetics (labels, title, ...) and saving mv_data (xarray.Dataset): The extracted multidimensional dataset dim (str, optional): The required parameter dimension of the 1d bifurcation diagram. dims (str, optional): The required parameter dimensions (x, y) of the 2d-bifurcation diagram. analysis_steps (Sequence): The analysis steps that are to be made until one is conclusive. Applied per universe. - If seq of str: The str will also be used as attractor key for plotting if the test is conclusive. - If seq of Tuple(str, str): The first str defines the attractor key for plotting, the second str is a key within custom_analysis_funcs. Default analysis_funcs are: - endpoint: utopya.dataprocessing.find_endpoint - fixpoint: utopya.dataprocessing.find_fixpoint - multistability: utdp.find_multistability - oscillation: utdp.find_oscillation - scatter: resolve_scatter custom_analysis_funcs (dict): A collection of custom analysis functions that will overwrite the default analysis funcs (recursive update). analysis_kwargs (dict, optional): The entries need to match the analysis_steps. The subentry (dict) is passed on to the analysis function. visualization_kwargs (dict, optional): The entries need to match the analysis_steps. The subentry (dict) is used to configure a rectangle to visualize the conclusive analysis step. Is passed to matplotlib.patches.rectangle. xy, width, height, and angle are ignored and set automatically. Required in 2d bifurcation diagram. to_plot (dict, optional): The configuration for the data to plot. The entries of this key need to match the data_vars selected in mv_data. It is used to visualize the state of the attractor additionally to the visualization kwargs. Only for 1d-bifurcation diagram. sub_keys: - ``label`` (str, optional): label in plot - ``plot_kwargs`` (dict, optional): passed to scatter for every universe - color (str, recommended): unique color for every data_variable accross universes **kwargs: Collection of optional dicts passed to different functions - plot_coords_kwargs (dict): Passed to ax.scatter to mark the universe's center in the bifurcation diagram - rectangle_map_kwargs (dict): Passed to utopya.eval.plots._utils.calc_pxmap_rectangles - legend_kwargs (dict): Passed to ax.legend """ def resolve_analysis_steps( analysis_steps: Sequence[Union[str, Tuple[str, str]]] ) -> Sequence[Tuple[str, str]]: """Resolve instance of str to Tuple[str, str] in sequence Args: analysis_steps (Sequence[Union[str, Tuple[str, str]]]): The original sequence Returns: analysis_steps (Sequence[Tuple[str, str]]): The sequence of attractor_key and analysis_func pairs. """ for i, analysis_step in enumerate(analysis_steps): # get key and func for the analysis step if isinstance(analysis_step, str): analysis_steps[i] = [analysis_step, analysis_step] return analysis_steps def resolve_to_plot_kwargs(to_plot: dict) -> dict: """Resolves the to_plot dict, e.g. adding labels if not explicitly specified. Args: to_plot (dict): The to_plot dict to parse Returns: dara_vars_plot_kwargs (dict): A dict with the 'plot_kwargs' for every data_var in to_plot. """ if not to_plot: return {} data_vars_plot_kwargs = {} for k, v in to_plot.items(): plot_kwargs = v.get("plot_kwargs", {}) if not plot_kwargs.get("label"): plot_kwargs["label"] = v.get("label", k) data_vars_plot_kwargs[k] = plot_kwargs return data_vars_plot_kwargs def create_legend_handles( *, visualization_kwargs: dict, data_vars_plot_kwargs: dict ): """Creates legend handles Processes entries in data_vars_plot_kwargs (from to_plot) and visualization_kwargs. Args: visualization_kwargs (dict): The visualization kwargs data_vars_plot_kwargs (dict): The resolved entries in to_plot Returns: Tuple[list, list]: Tuple of legend handles and legend labels lists as required by ax.legend(handles, labels) data_vars_plot_kwargs: The updated entries data_vars_plot_kwargs as required by plot_attractor. """ # Some defaults circle_kwargs = dict(xy=(0.5, 0.5), radius=0.25, edgecolor="none") rect_kwargs = dict( xy=(0.0, 0.0), height=0.75, width=1.0, edgecolor="none" ) # Lists to be populated for matplotlib legend legend_handles = [] legend_labels = [] for k, kwargs in data_vars_plot_kwargs.items(): label = kwargs.pop("label", kwargs) kwargs["linewidth"] = kwargs.get("linewidth", 0.0) data_vars_plot_kwargs[k] = kwargs # Determine color if "color" in kwargs: color = kwargs["color"] elif "cmap" in kwargs: cmap =["cmap"]) color = cmap(1.0) else: log.warning(f"No color defined for data_var '{k}'!") color = None # Create and add the handle and the label legend_handles.append(Circle(**circle_kwargs, facecolor=color)) legend_labels.append(label) for k, kwargs in visualization_kwargs.items(): if "to_plot" in kwargs: for dvar_name, dvar_kwargs in kwargs["to_plot"].items(): # Make sure a linewidth is set dvar_kwargs["linewidth"] = dvar_kwargs.get( "linewidth", 0.0 ) kwargs["to_plot"][dvar_name] = dvar_kwargs # Create and add the handle and the label legend_handles.append( Rectangle(**rect_kwargs, **dvar_kwargs) ) legend_labels.append(dvar_kwargs.get("label", dvar_name)) else: kwargs["linewidth"] = kwargs.get("linewidth", 0.0) data_vars_plot_kwargs[k] = kwargs legend_handles.append(Rectangle(**rect_kwargs, **kwargs)) legend_labels.append(kwargs.get("label", k)) return [legend_handles, legend_labels], data_vars_plot_kwargs def apply_analysis_steps( data: xr.Dataset, analysis_steps: Sequence[Union[str, Tuple[str, str]]], *, analysis_funcs: dict, analysis_kwargs: dict, ): """Perform the sequence of analysis steps until the first conclusive. Args: data (xarray.Dataset): The data to analyse. analysis_steps (Sequence[Union[str, Tuple[str, str]]]): The analysis steps that are to be made until one is conclusive. Applied per universe. analysis_funcs (dict): The entries need to match the analysis_steps. Map of the analysis_steps to their Callables analysis_kwargs (dict): The entries need to match the analysis_steps. The subentry (dict) is passed on to the corresponding analysis function. Returns: analysis_key (str): The key of the conclusive analysis step. attractor (xarray.Dataset): The data corresponding to this analysis. """ for analysis_key, analysis_func in analysis_steps: analysis_func_kwargs = analysis_kwargs.get(analysis_func, {}) # resolve the analysis function from its name if isinstance(analysis_func, str): if analysis_func in analysis_funcs: analysis_func = analysis_funcs[analysis_func] else: # Try to get it from dataprocessing ... might fail. analysis_func = getattr( utopya.dataprocessing, analysis_func ) # Perfom the analysis step conclusive, attractor = analysis_func(data, **analysis_func_kwargs) # Return if conclusive if conclusive: return analysis_key, attractor # Return non-conclusive return None, None def resolve_rectangle(coord: dict, rectangles: xr.Dataset) -> Rectangle: """Resolve the rectangle patch at this coordinate Args: coord (dict): The bifurcation parameter's coordinate rectangles (xarray.Dataset): The rectangles that cover the 2D space spanned by the coordiantes. The `coord` should be one entry of rectangles.coords. Raises: ValueError: Coordinate not available in rectangles. Returns: Rectangle: A rectangle around a universe with coord and shape defined by rectangles. """ try: rectangle = rectangles.sel(coord) except Exception as exc: raise ValueError( "The requested paramspace coordinate(s) {} are " "not coordinates of rectangles {}. Plot failed." "".format(coord, rectangles.coords) ) from exc rect_spec = rectangle["rect_spec"] return Rectangle(*rect_spec.item()) def append_vis_patch( attrator_key: str, attractor: xr.Dataset, vis_patches: dict, vis_kwargs: dict, **resolve_rectangle_args, ): """Append visualization patch Performs postprocess for - attractor key 'fixpoint' and 'endpoint' if: 'to_plot' in entry of vis_kwargs. Then finds the data_var with highest valued datapoint. Args: attractor_key (str): Key according to which to decode the attractor attractor (xarray.Dataset): The Dataset with the encoded attractor information. See possible encodings vis_patches (dict): The map of attractor_key to List[Rectangle] where to append the new patch vis_kwargs (dict): The visualization kwargs **resolve_rectangle_args: Passed on to resolve_rectangle Raises: ValueError: Bad postprocess key Returns: vis_patches (dict): The new map of attractor_key to List[Rectangle] Deleted Parameters: resolve_rectangle_args (dict): Args as required by resolve_rectangle """ # Depending on the kind of attractor, add different patches kwargs = vis_kwargs.get(attractor_key) if kwargs is None: return vis_patches # Postprocess fixpoint and to_plot, append rectangle if ( attractor_key == "fixpoint" or attractor_key == "endpoint" ) and "to_plot" in kwargs: max_value = -np.inf for data_var_name, data_var in attractor.data_vars.items(): if data_var.max() > max_value: max_value = data_var.max() max_name = data_var_name rect = resolve_rectangle(**resolve_rectangle_args) attractor_var_key = attractor_key + "_" + max_name vis_patches[attractor_var_key].append(rect) # Append rectangle elif vis_patches.get(attractor_key) is not None: rect = resolve_rectangle(coord=coord, rectangles=rects) vis_patches[attractor_key].append(rect) return vis_patches def append_plot_attractor( attractor_key: str, attractor: xr.Dataset, *, coord: float = None, scatter_kwargs: list, **plot_kwargs, ): """Resolves how to plot attractor of specified type at specific bifurcation parameter value. Args: attractor_key (str): Key according to which to decode the attractor attractor (xarray.Dataset): The Dataset with the encoded attractor information. See possible encodings coord (float, optional): The bifurcation parameter's coordinate, if None its derived from the attractors coordinates scatter_kwargs (list): The list of scatter datasets where to append the new scatter. plot_kwargs (dict, optional): The kwargs used to specify ax.scatter where the entries match the attractor.data_vars Possible encodings, i.e. values for ``attractor_key`` (of type :py:class:`xarray.Dataset` or a compatible type): - ``fixpoint``: dataset with dimensions ``()`` - ``scatter``: dataset with dimensions ``(time: >=1)`` - ``multistability``: dataset with dimensions ``(<initial_state>: >= 1)`` - ``oscillation``: dataset with dimensions ``(osc: 2)``, the minimum and maximum .. note:: The attractor must contain the bifurcation parameter coordinate Raises: KeyError: Unknown ``attractor_key`` or no bifurcation coordinate ValueError: Attractor encoding mismatched with the given ``attractor_key`` Returns: dict: The new list of scatter datasets """ # Get the bifurcation parameter coordinate if not coord: try: coord = attractor[dim] except KeyError as err: raise ValueError( f"No bifurcation parameter coordinate '{dim}' " "could be found! Either have it as a " "coordinate in 'attractor' or pass it to " "'plot_attractor' explicitly." ) from err # Resolve the scatter kwargs depending on attractor key if attractor_key in ("fixpoint", "endpoint", "multistability"): for data_var_name, data_var in attractor.data_vars.items(): data_var = data_var.where(data_var != np.nan, drop=True) entries = 1 if data_var.shape: entries = len(data_var) scatter_kwargs.append( dict( x=[coord] * entries, y=data_var, **plot_kwargs.get(data_var_name, {}), ) ) elif attractor_key == "scatter": for data_var_name, data_var in attractor.data_vars.items(): if "cmap" in plot_kwargs.get(data_var_name, {}): scatter_kwargs.append( dict( x=[coord] * len(, y=data_var, c=attractor["time"], **plot_kwargs.get(data_var_name, {}), ) ) else: scatter_kwargs.append( dict( x=[coord] * len(, y=data_var, **plot_kwargs.get(data_var_name, {}), ) ) elif attractor_key == "oscillation": for data_var_name, data_var in attractor.data_vars.items(): scatter_kwargs.append( dict( x=[coord] * len(, y=data_var, **plot_kwargs.get(data_var_name, {}), ) ) elif attractor_key: raise ValueError( f"Invalid attractor-key '{attrator_key}'! " "Available keys: 'endpoint', fixpoint'," " 'multistability', 'scatter', 'oscillation'." ) return scatter_kwargs def resolve_scatter( data: xr.Dataset, *, spin_up_time: int = 0, **kwargs ) -> tuple: """A mock analysis function to plot all times larger than a spin up time. Args: data (xarray.Dataset): The dataset to analyse spin_up_time (int, optional): The spin-up-time **kwargs: Ignored """ return True, data.where(data.time >= spin_up_time, drop=True) # ......................................................................... # Check argument values if not dim and not dims: raise ValueError( "No dim (str) or dims (Tuple[str, str]) specified. " "Use dim for a 1d-bifurcation diagram and dims for a " "2d-bifurcation diagram." ) if dim and dims: raise ValueError( "dim='{}' and dims='{}' specified. " "Use either dim for a 1d-bifurcation diagram or dims " "for a 2d-bifurcation diagram." "".format(dim, dims) ) if dims is not None and len(dims) != 2: raise ValueError( f"Argument dims should be of length 2, but was: {dims}" ) # TODO In the future, consider not using `dim` below here but handling it # via the length of `dims`. # Default values if visualization_kwargs is None: visualization_kwargs = {} if analysis_kwargs is None: analysis_kwargs = {} # Resolve legend handles and visualization kwargs data_vars_plot_kwargs = resolve_to_plot_kwargs(to_plot) legend_handles, data_vars_plot_kwargs = create_legend_handles( visualization_kwargs=visualization_kwargs, data_vars_plot_kwargs=data_vars_plot_kwargs, ) # Define default analysis functions analysis_funcs = dict( endpoint=utdp.find_endpoint, fixpoint=utdp.find_fixpoint, multistability=utdp.find_multistability, oscillation=utdp.find_oscillation, scatter=resolve_scatter, ) # If given, update if custom_analysis_funcs: log.debug("Updating with custom analysis functions ...") analysis_funcs = recursive_update( analysis_funcs, custom_analysis_funcs ) analysis_steps = resolve_analysis_steps(analysis_steps) # Obtain the rectangles covering space spanned by the coordinates rectangle_map_kwargs = kwargs.get("rectangle_map_kwargs", {}) if dim: rects, limits = calc_pxmap_rectangles( x_coords=mv_data[dim].values, y_coords=None, **rectangle_map_kwargs ) elif dims: rects, limits = calc_pxmap_rectangles( x_coords=mv_data[dims[0]].values, y_coords=mv_data[dims[1]].values, **rectangle_map_kwargs, ) # Obtain the list of param_coords to iterate if dim: param_iter = mv_data[dim].values elif dims: param_iter = itertools.product( mv_data[dims[0]].values, mv_data[dims[1]].values ) # Map of analysis_key to list[mpatch.Rectangle] vis_patches = {} for analysis_key, _ in analysis_steps: if not visualization_kwargs.get(analysis_key): continue if "to_plot" in visualization_kwargs[analysis_key]: for var_key, _ in visualization_kwargs[analysis_key][ "to_plot" ].items(): analysis_var_key = analysis_key + "_" + var_key vis_patches[analysis_var_key] = [] else: vis_patches[analysis_key] = [] # The List[dict] passed to ax.scatter scatter_kwargs = [] scatter_coords_kwargs = [] # Iterate the parameter coordinates for param_coord in param_iter: # Resolve the param_coord to dict if dim: param_coord = {dim: param_coord} elif dims: param_coord = {dims[0]: param_coord[0], dims[1]: param_coord[1]} # Plot coord if dim and kwargs.get("plot_coords_kwargs"): plot_coords_kwargs = kwargs.get("plot_coords_kwargs") scatter_coords_kwargs.append( { "x": param_coord[dim], "y": plot_coords_kwargs.pop("y", 0.0), **plot_coords_kwargs, } ) if dims and kwargs.get("plot_coords_kwargs"): scatter_coords_kwargs.append( { "x": param_coord[dims[0]], "y": param_coord[dims[1]], **kwargs.get("plot_coords_kwargs"), } ) # Select the data and analyse data = mv_data.sel(param_coord) attractor_key, attractor = apply_analysis_steps( data, analysis_steps, analysis_funcs=analysis_funcs, analysis_kwargs=analysis_kwargs, ) # If conclusive, append a rectangular patch to the attractor_key's # patch collection if attractor_key: # Determine coordinate value if dim: rect_map_kwargs = kwargs.get("rectangle_map_kwargs", {}) y = rect_map_kwargs.get("default_pos", (0.0, 0.0))[1] coord = dict(x=param_coord[dim], y=y) elif dims: coord = dict(x=param_coord[dims[0]], y=param_coord[dims[1]]) vis_patches = append_vis_patch( attractor_key, attractor, vis_patches, visualization_kwargs, coord=coord, rectangles=rects, ) # For 1d case ... if dim and to_plot: scatter_kwargs = append_plot_attractor( attractor_key, attractor, coord=param_coord[dim], scatter_kwargs=scatter_kwargs, **data_vars_plot_kwargs, ) # Draw collection of visualization patches for analysis_key, _ in analysis_steps: if not visualization_kwargs.get(analysis_key): continue if analysis_key in visualization_kwargs: vis_kwargs = visualization_kwargs[analysis_key] if "to_plot" in vis_kwargs: for var_key, var_kwargs in vis_kwargs["to_plot"].items(): attractor_var_key = analysis_key + "_" + var_key pc = PatchCollection( vis_patches[attractor_var_key], **var_kwargs ) else: pc = PatchCollection(vis_patches[analysis_key], **vis_kwargs) # else: nothing to do # Scatter the universe's coordinates for kws in scatter_coords_kwargs:**kws) # Scatter the attractor for kws in scatter_kwargs:**kws) # Provide PlotHelper defaults hlpr.provide_defaults("set_limits", **limits) if dim: hlpr.provide_defaults("set_labels", x=dim, y="state") elif dims: hlpr.provide_defaults("set_labels", x=dims[0], y=dims[1]) if legend_handles: legend_handles[0], legend_handles[1], handler_map={Circle: HandlerEllipse()}, **kwargs.get("legend_kwargs", {}), )