code: # inferred_model.py contains functions to visualize
  # timeseries, models for predicting them, and their comparison.
  #
  # Authors:
  # Charlie Brummitt <charles_brummitt@hms.harvard.edu>, Github: cbrummitt
  # Andres Gomez <Andres_Gomez@hks.harvard.edu>
  import itertools
  
  import matplotlib as mpl
  import matplotlib.pyplot as plt
  import numpy as np
  import scipy
  from matplotlib.patches import Ellipse
  from mpl_toolkits.axes_grid1 import make_axes_locatable
  from scipy.spatial import ConvexHull
  from sklearn.kernel_ridge import KernelRidge
  from sklearn.metrics import mean_squared_error
  from sklearn.utils import check_array
  
  from palettable.colorbrewer.qualitative import Dark2_3
  
  from .utils import (convert_None_to_empty_dict_else_copy, create_fig_ax,
                      make_color_cycler, maybe_save_fig)
  
  country_color_dict = dict(zip([
      'ARG', 'AUS', 'BRA', 'CAN', 'CHN', 'COL', 'ESP', 'FRA', 'GBR', 'GRC',
      'HND', 'IDN', 'IND', 'JPN', 'KEN', 'KOR', 'MDG', 'MEX', 'RUS', 'THA',
      'USA', 'BEL', 'NLD'],
      make_color_cycler()))
  
  
  def make_data_for_streamplot(
          data, model, n_points_each_dim, model_predicts_change=True,
          mask_velocity_to_convex_hull_of_data=False,
          dim_reducer=None,
          dimensions_to_keep=(0, 1), aggregator='mean'):
      """Make data for a stream plot or quiver plot by creating a grid of equally
      spaced points that contains the data, predicting at those grid points using
      the model, and returning those grid points as 1D vectors and the
      predictions as tensors.
  
      Parameters
      ----------
      data : array-like, shape [n_samples, n_features]
          The data used to train the model, where `model` was fit with
          `model.fit(data, target)`.
  
      model : fitted estimator with a predict method
  
      n_points_each_dim : int or array of ints of length equal to data.ndim
          The number of grid points to use in each dimension
  
      model_predicts_change : bool, optional, default: True
          Whether the model predicts the change in the system from one time step
          to the next (model_predicts_change == True) or predicts the value of
          the system at the next time step.
  
      mask_velocity_to_convex_hull_of_data : bool, optional, default: False
          If True, then mask the `velocities` arrays by excluding points outside
          the convex hull of `data`.
  
      dimensions_to_keep : int, optional, default: 2
          Which dimensions to keep in the tensors of grid points and
          predictions. The other dimensions are aggregated over.
  
      aggregator : {'mean', 'median'} or user-defined function
          What function to use to aggregate over the higher-order dimensions in
          order to reduce the dimensions of the tensor to keep only those
          dimensions in the tuple `dimensions_to_keep`.
          The function must take as input a tensor (of shape specified by
          `n_points_each_dim`) and an `axis` argument that, when given a tuple
          of integers for which axes to aggregate over, reduces the dimension of
          the tensor by applying the function to those higher-order axes.
  
      Returns
      -------
      grid_points_1d_for_each_dim : `len(dimensions_to_keep)` many arrays of
      length specified by n_points_each_dim
          The locations of the grid points of the grid at which the model
          predicts the change in the system.
  
      predictions : arrays of shape specified by the dimensions in the tuple
      dimensions_to_keep with dimensions specified by `n_points_each_dim`
          The predictions at the grid, with the higher-order dimensions reduced
          according to the `aggregator` function.
  
      Examples
      --------
      >>> grid_points, predictions = make_data_for_streamplot(data, model, 20)
      >>> plt.streamplot(*grid_points, *predictions)
      """
      data = check_array(data)
      grid_points, grids = bounding_grid(data, n_points_each_dim)
  
      velocities = predict_velocities_on_grid(
          model, grids, model_predicts_change=model_predicts_change)
  
      if mask_velocity_to_convex_hull_of_data:
          velocities = mask_arrays_with_convex_hull(
              velocities, grid_points, ConvexHull(data))
  
      def reduce_dimensions_by_aggregating_over_axes(tensors):
          return reduce_dimensions_of_grids(
              tensors,
              dimensions_to_keep=dimensions_to_keep, aggregator=aggregator)
  
      grids = reduce_dimensions_by_aggregating_over_axes(grids)
      velocities = reduce_dimensions_by_aggregating_over_axes(velocities)
  
      grid_points_1d_for_each_dim = [grid_points[i] for i in dimensions_to_keep]
      # Must transpose the velocities for streamplot and quiver plot
      velocities = [v.T for v in velocities]
  
      return grid_points_1d_for_each_dim, velocities
  
  
  def mask_arrays_with_convex_hull(arrays, grid_points, convex_hull):
      """Mask the arrays to retain only those in the given convex hull.
  
      Parameters
      ----------
      arrays : list of tensors of shapes n1, n2, n3, ...
  
      grid_points : list of 1D arrays of length n1, n2, n3, ...
          The arrays are data measured on the Cartesian product of the
          `grid_points`
  
      convex_hull : ConvexHull of the data
  
      Returns
      -------
      arrays_masked : NumPy masked arrays of shapes n1, n2, n3, ...
          The given `arrays` with points outside `convex_hull` masked.
      """
      mask = np.zeros(arrays[0].shape, dtype='bool').flatten()
      for i, point in enumerate(itertools.product(*grid_points)):
          mask[i] = not point_in_hull(point, convex_hull)
      mask = mask.reshape(arrays[0].shape)
      return [np.ma.masked_array(data=arr, mask=mask) for arr in arrays]
  
  
  def point_in_hull(point, hull, tolerance=1e-12):
      return all(
          (np.dot(eq[:-1], point) + eq[-1] <= tolerance)
          for eq in hull.equations)
  
  
  def example_make_data_for_streamplot(n_samples=10, n_features=2, p=4):
      fake_data = np.random.randn(n_samples, n_features)
      fake_target = np.random.randn(n_samples, n_features)
      krr = KernelRidge()
      krr.fit(fake_data, fake_target)
  
      grid_points, predictions = make_data_for_streamplot(
          fake_data, krr, p, model_predicts_change=False)
      return grid_points, predictions
  
  
  def bounding_grid(data, n_points_each_dim=30, max_num_points=1e7):
      """Creates an evenly spaced grid that contains the data. It returns a
      numpy meshgrid with indexing 'ij'.
  
      Parameters
      ----------
      data : array-like, shape [num_samples, num_features] with num_features >= 2
          The data to create a grid around.
  
      n_points_each_dim : scalar or tuple of length data.shape[1]
          The number of grid points to use each dimension. If this parameter
          is a scalar, then it is duplicated for every column of `data`.
  
      max_num_points : int
          The maximum number of points allowed in the grid; above this value,
          a ValueError is raised.
  
      Returns
      -------
      grid_points : list of vectors of lengths given by n_points_each_dim
          The grid points in dimensions 0, 1, ..., data.shape[1]
  
      meshgrid : tensors of grid locations
          The output of np.meshgrid. data.shape[1] many tensors are returned.
  
      Raises
      ------
      ValueError if `n_points_each_dim` is a tuple or list of length different
      from the number of columns in `data`
  
      ValueError if the number of points in the grid is larger than
      `max_num_points`
  
      Examples
      --------
      >>> bounding_grid(np.arange(6).reshape(3, 2), 3)
      [array([[ 0.,  0.,  0.],
              [ 2.,  2.,  2.],
              [ 4.,  4.,  4.]]),
       array([[ 1.,  3.,  5.],
              [ 1.,  3.,  5.],
              [ 1.,  3.,  5.]])]
      """
      dim_n_points_each_dim = len(np.atleast_1d(n_points_each_dim))
      if dim_n_points_each_dim > 1 and dim_n_points_each_dim != data.shape[1]:
          raise ValueError('The number of points requested for each dimension '
                           'is a vector of length {}, but the data has {} col'
                           'umns'.format(dim_n_points_each_dim, data.shape[1]))
      min_of_columns = np.floor(np.min(data, axis=0))
      max_of_columns = np.ceil(np.max(data, axis=0))
  
      n_points_each_dim = duplicate_to_desired_length_if_scalar(
          n_points_each_dim, data.shape[1])
  
      n_points_total = np.prod(n_points_each_dim)
      if n_points_total > max_num_points:
          raise ValueError(
              'The bounding grid has {} many points, which exceeds the maximum'
              '`max_num_points`={}'.format(n_points_total, max_num_points))
  
      start_stop_num = zip(min_of_columns, max_of_columns, n_points_each_dim)
      grid_points = [np.linspace(lo, hi, num) for lo, hi, num in start_stop_num]
      meshgrid = np.meshgrid(*grid_points, indexing='ij')
      return grid_points, meshgrid
  
  
  def duplicate_to_desired_length_if_scalar(x, desired_length):
      x = np.atleast_1d(x)
      if len(x) == 1 and desired_length > 1:
          x = np.repeat(x, desired_length)
      return x
  
  
  def predict_velocities_on_grid(model, meshgrids, model_predicts_change=True):
      """Predict the velocities of a system at a mesh of grid points, and return
      the grid and the predicted velocities, which can then be provided as input
      to streamplot and quiver.
      """
      meshgrids_long_format = np.array([ary.flatten() for ary in meshgrids]).T
  
      n_features = len(meshgrids)
      n_points_each_dim = meshgrids[0].shape
      n_grid_points = np.prod(n_points_each_dim)
      assert meshgrids_long_format.shape == (n_grid_points, n_features)
  
      predictions_long_format = model.predict(
          meshgrids_long_format)
      assert predictions_long_format.shape == (n_grid_points, n_features)
  
      predictions_shape_grids = [
          predictions_long_format[:, i].reshape(*n_points_each_dim)
          for i in range(n_features)]
      assert len(predictions_shape_grids) == n_features
      assert predictions_shape_grids[0].shape == meshgrids[0].shape
  
      if model_predicts_change:
          velocities = predictions_shape_grids
      else:
          meshgrids_preds = zip(meshgrids, predictions_shape_grids)
          velocities = ([pred - grid for grid, pred in meshgrids_preds])
      return velocities
  
  
  def aggregate_dimensions_of_grid_points_and_velocities(
          grid_points, velocities, dimensions_to_keep, aggregator='mean'):
      """Select dimensions and aggregate over other dimensions.
  
      Parameters
      ----------
      grid_points : list of 1D arrays
  
      velocities : list of ND arrays of shape given by the lengths of the
      elements of grid_points
  
      dimensions_to_keep : tuple of int's length 2
          Which dimensions (features) to keep. Each entry in the tuple is an
          int between 0 and n_features - 1 (inclusive).
  
      aggregator : {'mean', 'median', or callable}, default: 'mean'
          How to aggregate over axes of the tensor. If callable, it must take
          as input the tensor and a keyword argument `axis` that is given a
          tuple of the indices of the axes to aggregate over.
      """
      grid_points = [grid_points[dim] for dim in dimensions_to_keep]
      velocities = reduce_dimensions_of_grids(
          velocities, dimensions_to_keep, aggregator=aggregator)
      return grid_points, velocities
  
  
  def reduce_dimensions_of_grids(
          grids, dimensions_to_keep=(0, 1), aggregator='mean'):
      """Reduce the dimensions of tensors by aggregating over some dimensions.
  
      Inputs
      ------
      grids : a list of arrays of the same shape
  
      dimensions_to_keep : tuple of int's between 0 and the number of dimensions
      of the tensors in `grids`
          Which dimensions to keep. The other dimensions are aggregated over.
  
      aggregator : {'mean', 'median', or callable}, default: 'mean'
          How to aggregate over axes of the tensor. If callable, it must take
          as input the tensor and a keyword argument `axis` that is given a tuple
          of the indices of the axes to aggregate over.
      """
      for i in range(len(grids)):
          for j in range(i + 1, len(grids)):
              assert grids[i].shape == grids[j].shape
  
      axes_to_aggregate = tuple(
          dim for dim in range(grids[0].ndim) if dim not in dimensions_to_keep)
      aggregation_function = {
          'mean': np.mean, 'median': np.median}.get(aggregator, aggregator)
      return [aggregation_function(grids[i], axis=axes_to_aggregate)
              for i in dimensions_to_keep]
  
  
  def quiver_plot_empirical(
          data, target, dimensions_to_keep=(0, 1),
          color_values='speed', colorbar_label='speed',
          ax=None, save_fig=None,
          xlabel='first component', ylabel='second component', **subplots_kws):
      """Create a quiver plot of a dynamical system.
  
      Parameters
      ----------
      data : array-like, shape [n_samples, n_features]
          Starting points of the arrows. If n_features > 3, then only the first
          two features are plotted.
  
      target : array-like, shape [n_samples, n_features]
          Ending points of the arrows. If n_features > 3, then only the first
          two features are plotted.
  
      dimensions_to_keep : tuple of int's length 2
          Which dimensions (features) to plot. Each entry in the tuple is an
          int between 0 and n_features - 1 (inclusive).
  
      color_values : string or 2D numpy array
          Data for the colors of the arrows in the streamplot. If color_values is
          'speed', then color_values is the magnitude of the velocity.
  
      colorbar_label : str, optional, default: 'speed'
          The label of the color bar
  
      ax : matplotlib axis, optional, default: None
          The axis on which to draw the plot. If None, then an axis is created.
  
      xlabel, ylabel : str, optional,
      default: 'first component', 'second component'
          The labels of the axes
  
      subplots_kws : keyword arguments to pass to plt.subplots, default: None
  
      Returns
      -------
      fig, ax : matplotlib Figure, Axis
      """
      assert len(dimensions_to_keep) == 2
      for dim in dimensions_to_keep:
          assert 0 <= dim <= data.shape[1]
  
      data = check_array(data)[:, dimensions_to_keep]
      target = check_array(target)[:, dimensions_to_keep]
      X, Y, U, V = make_X_Y_U_V_for_quiver_plot(data, target)
      return quiver_plot(X, Y, U, V, color_values=color_values,
                         colorbar_label=colorbar_label, ax=ax, save_fig=save_fig,
                         xlabel=xlabel, ylabel=ylabel, **subplots_kws)
  
  
  def make_X_Y_U_V_for_quiver_plot(data, target, target_is_difference=False):
      data, target = check_array(data), check_array(target)
      X, Y = data.T
      if target_is_difference:
          U, V = target.T
      else:
          U, V = (target - data).T
      return X, Y, U, V
  
  
  def quiver_plot(
          X, Y, U, V,
          show_colorbar=True, color_values='speed', colorbar_label='speed',
          xlabel='first component', ylabel='second component',
          ax=None, save_fig=None,
          subplots_kws=None,
          plot_kws={'cmap': mpl.cm.viridis,
                    'angles': 'xy', 'scale': 1, 'scale_units': 'xy'}):
      """Create a quiver plot of a dynamical system with units given by the
      velocities.
  
      Parameters
      ----------
      X, Y : vectors of length equal to the number of samples
          Horizontal and vertical locations of the tails of the arrows
  
      U, V : vectors of length equal to the number of samples
          The velocities at the corresponding points in `X`, `Y`
  
      show_colorbar : bool, default: True
          Whether to show a colorbar.
  
      color_values : string or 2D numpy array
          Data for the colors of the arrows in the streamplot. If color_values is
          'speed', then color_values is the magnitude of the velocity.
  
      colorbar_label : str, optional, default: 'speed'
          The label of the color bar
  
      xlabel, ylabel : str, optional, default: 'first component',
      'second component'
          The labels of the axes
  
      ax : matplotlib axis, optional, default: None
          The axis on which to draw the plot. If None, then an axis is created.
  
      save_fig : str or None, default: None
          If not None, then save the figure to the path specified by this string.
  
      subplots_kws : None or dict
          Keyword arguments to pass to plt.subplots
  
      plot_kws : None or dict
          Keyword arguments to pass to plt.quiver
  
      Returns
      -------
      fig, ax : matplotlib Figure, Axis
      """
      subplots_kws = convert_None_to_empty_dict_else_copy(subplots_kws)
      quiver_kws = convert_None_to_empty_dict_else_copy(plot_kws)
  
      fig, ax = create_fig_ax(ax, **subplots_kws)
      if color_values == 'speed':
          color_values = np.sqrt(U**2 + V**2)
      quiv = ax.quiver(
          X, Y, U, V, color_values, **quiver_kws)
      if show_colorbar:
          divider = make_axes_locatable(ax)
          cax = divider.append_axes('right', size='4%', pad=0.04)
          cbar = fig.colorbar(mappable=quiv, cax=cax, orientation='vertical')
          cbar.ax.set_ylabel(colorbar_label)
      ax.set_xlabel(xlabel)
      ax.set_ylabel(ylabel)
      maybe_save_fig(fig, save_fig)
      return fig, ax
  
  
  def stream_plot(
          X, Y, U, V,
          show_colorbar=True, color_values='speed', colorbar_label='speed',
          subplots_kws=None, plot_kws=None,
          ax=None, save_fig=None,
          xlabel='first component', ylabel='second component'):
      """Create a stream plot of a dynamical system with units given by the
      velocities.
  
      Parameters
      ----------
      X, Y : vectors of length equal to the number of samples
          Horizontal and vertical locations of the tails of the arrows
  
      U, V : vectors of length equal to the number of samples
          The velocities at the corresponding points in `X`, `Y`. These can be
          masked arrays; their missing values will be filled with zero.
  
      show_colorbar : bool, default: True
          Whether to show a colorbar.
  
      color_values : string or 2D numpy array
          Data for the colors of the arrows in the streamplot. If color_values is
          'speed', then color_values is the magnitude of the velocity.
  
      colorbar_label : str, optional, default: 'speed'
          The label of the color bar
  
      xlabel, ylabel : str, optional, default: 'first component',
      'second component'
          The labels of the axes
  
      ax : matplotlib axis, optional, default: None
          The axis on which to draw the plot. If None, then an axis is created.
  
      save_fig : str or None, default: None
          If not None, then save the figure to the path specified by this string.
  
      subplots_kws : keyword arguments to pass to plt.subplots
  
      plot_kws : keyword arguments to pass to plt.subplots
  
      Returns
      -------
      fig, ax : matplotlib Figure, Axis
      """
      subplots_kws = convert_None_to_empty_dict_else_copy(subplots_kws)
      plot_kws = {
          **{'cmap': mpl.cm.viridis},
          **convert_None_to_empty_dict_else_copy(plot_kws)}
  
      U = fill_mask_if_needed(U, fill_value=0)
      V = fill_mask_if_needed(V, fill_value=0)
  
      fig, ax = create_fig_ax(ax, **subplots_kws)
      if color_values == 'speed':
          color_values = np.sqrt(U**2 + V**2)
      strm = ax.streamplot(X, Y, U, V, color=color_values, **plot_kws)
      ax.set_xlabel(xlabel)
      ax.set_ylabel(ylabel)
      if show_colorbar:
          divider = make_axes_locatable(ax)
          cax = divider.append_axes('right', size='4%', pad=0.04)
          cbar = fig.colorbar(mappable=strm.lines, cax=cax,
                              orientation='vertical')
          cbar.ax.set_ylabel(colorbar_label)
      maybe_save_fig(fig, save_fig)
      return fig, ax
  
  
  def fill_mask_if_needed(x, fill_value=0):
      if hasattr(x, 'filled'):
          return x.filled(fill_value=fill_value)
      else:
          return x
  
  
  def plot_ellipse(ax, position, dimensions, label='kernel scale:'):
      width, height = duplicate_to_desired_length_if_scalar(dimensions, 2)
      ellipse = Ellipse(position, width, height,
                        color='y', alpha=.5, clip_on=True)
      ax.add_artist(ellipse)
      ax.text(position[0] - width / 4., position[1] + height / 4.,
              '{lab}\n{w:.2f} x {h:.2f}'.format(lab=label, w=width, h=height))
  
  
  def select_trajectories(panel, items=None):
      """Select trajectories of certain items in a panel.
  
      Parameters
      ----------
      panel : pandas Panel
          A panel of trajectories, with time on the major axis, items denoting
          different trajectories, and features on the minor axis.
  
      items : iterable of strings (each string an item of the `panel`) or None
          The items of the panel to select. If None, then use all items in the
          panel.
  
      Returns
      -------
      items_to_trajectories : dict
          Dictionary mapping the items to their trajectories.
      """
      items_to_trajectories = {}
      if items is None:
          items = panel.items
      for item in items:
          items_to_trajectories[item] = panel.loc[item].dropna(how='all').values
      return items_to_trajectories
  
  
  def iterated_predictions(
          panel, model, items=None, model_predicts_change=True,
          num_time_steps=100, index_of_initial_condition=-1):
      """Compute iterated predictions of certain items in a panel.
  
      Parameters
      ----------
      panel : pandas Panel
          A panel of trajectories, with time on the major axis, items denoting
          different trajectories, and features on the minor axis.
  
      items : list of strings (items in the panel) or None, default: None
          The items to select from the panel and to make predictions
  
      model : a fitted estimator that can predict the next year
  
      model_predicts_change : bool, optional, default: True
          Whether the model predicts the change in the system from one time step
          to the next (model_predicts_change == True) or predicts the value of
          the system at the next time step.
  
      num_time_steps : int or 'length_trajectory'
          The number of time steps to predict into the future. If num_time_steps
          is 'length_trajectory', then num_time_steps is set to the length of
          the trajectory of that time.
  
      index_of_initial_condition : int, optional, default: -1
          The index of the item's trajectory to use as initial conditon. If -1,
          then the initial condition is the last observation; if 0, then the
          intial condition is the first observation.
  
      Returns
      -------
      items_to_trajectories : dict mapping strings to arrays of shape
      [n_time_steps, n_features]
          Dictionary mapping the items to their trajectories.
      """
      items_to_trajectories = {}
      if items is None:
          items = panel.items
  
      for item in items:
          item_df = panel.loc[item].dropna(how='all')
          initial_condition = item_df.iloc[index_of_initial_condition].values
          if num_time_steps == 'length_trajectory':
              n_steps_to_predict = len(item_df)
          else:
              n_steps_to_predict = num_time_steps
          trajectory = np.empty((n_steps_to_predict, initial_condition.shape[0]))
          trajectory[0] = initial_condition
          for i in range(1, n_steps_to_predict):
              trajectory[i] = model.predict(trajectory[i - 1].reshape(1, -1))
              if model_predicts_change:
                  trajectory[i] += trajectory[i - 1]
          items_to_trajectories[item] = trajectory
      return items_to_trajectories
  
  
  def plot_item_trajectories(
          items_to_trajectories,
          items_to_colors=None,
          dimensions_to_keep=(0, 1),
          ax=None, text_labels=False, shift_text_labels=None,
          arrow_start_index=-2, subplots_kws=None, plot_kws=None,
          arrow_kws={'head_length': 0.5, 'head_width': 0.5},
          text_kws={'fontsize': 10, 'alpha': 1,
                    'bbox': {'alpha': 0.2, 'pad': 1}},
          use_plot_alpha_for_arrowhead=False,
          xlabel=None, ylabel=None,
          save_fig=None):
      """Plot iterated predictions of certain items in a panel.
  
      Parameters
      ----------
      items_to_trajectories : dict mapping strings to arrays of shape
      [n_time_steps, n_features], or pandas Panel
          Dictionary mapping the items to their trajectories. If a pandas Panel,
          then it is converted to a dictionary mapping items to trajectories.
          If `n_features` is > 2, then only the first two features are plotted.
  
      items_to_colors : dict mapping items to colors, or None
          The keys of this dictionary are which items' trajectories should be
          plotted. If None, then default colors are used.
  
      ax : matplotlib axis, optional, default: None
          The axis on which to draw the plot. If None, then an axis is created.
  
      text_labels : bool, optional, default: False
          Whether to show text labels of the items
  
      shift_text_labels : dict, optional, default: None
          Provide a dictionary to shift the text labels. By default the labels
          appear at the third-to-last point in the trajectory.
  
      arrow_start_index : int, optional, default: -3
          The tail of the arrow is located at the trajectory sliced at this
          index. It must be between 0 and the second-to-last index; if not, then
          the second-to-last index is used.
  
      subplots_kws, plot_kws, arrow_kws, text_kws : dict, optional
          Keyword arguments passed to `subplots`, `plot`, `arrow`, `text`,
          respectively.
  
      use_plot_alpha_for_arrowhead : bool, default: False
          If True, and if no alpha is specified in `arrow_kws` and 'alpha' is
          specified in `plot_kws`, then set the opacity 'alpha' of the arrowhead
          to the value of alpha in `plot_kws`.
  
      xlabel, ylabel : None or str
          If not None, then use these as the labels of the x and y axes
  
      save_fig : str or None, default: None
          If not None, then save the figure to the path specified by this string.
      """
      items_to_colors = convert_None_to_empty_dict_else_copy(items_to_colors)
      subplots_kws = convert_None_to_empty_dict_else_copy(subplots_kws)
      plot_kws = convert_None_to_empty_dict_else_copy(plot_kws)
      arrow_kws = convert_None_to_empty_dict_else_copy(arrow_kws)
      text_kws = convert_None_to_empty_dict_else_copy(text_kws)
      shift_text_labels = convert_None_to_empty_dict_else_copy(shift_text_labels)
  
      if hasattr(items_to_trajectories, 'major_axis'):  # if it's a panel
          items_to_trajectories = select_trajectories(
              items_to_trajectories)
  
      color_cycler = make_color_cycler()
      fig, ax = create_fig_ax(ax, **subplots_kws)
  
      if xlabel:
          ax.set_xlabel(xlabel)
      if ylabel:
          ax.set_ylabel(ylabel)
  
      for item in items_to_trajectories:
          color = items_to_colors.get(item, next(color_cycler))
  
          trajectory = items_to_trajectories[item][:, dimensions_to_keep]
          plot_keyword_args = plot_kws.copy()
          plot_keyword_args.update(color=color)
          ax.plot(*trajectory.T, **plot_keyword_args)
  
          try:
              trajectory[arrow_start_index], trajectory[arrow_start_index + 1]
          except IndexError:
              arrow_start_index = len(trajectory) - 1
          if arrow_start_index == -1:
              print('trajectry for item ', item, ' is empty')
          else:
              x, y = trajectory[arrow_start_index]
              dx, dy = (trajectory[arrow_start_index + 1] -
                        trajectory[arrow_start_index])
              if use_plot_alpha_for_arrowhead:
                  if 'alpha' not in arrow_kws and 'alpha' in plot_kws:
                      arrow_kws['alpha'] = plot_kws['alpha']
              ax.arrow(x, y, dx, dy, color=color, **arrow_kws)
  
              if text_labels:
                  text_kws.update(color=color)
                  text_kws['bbox'].update(color=color)
                  x_text, y_text = (
                      trajectory[arrow_start_index] +
                      shift_text_labels.get(item, np.array([0, 0])))
                  ax.text(x_text, y_text, item, **text_kws)
  
      maybe_save_fig(fig, save_fig)
      return fig, ax
  
  
  def plot_model_predictions_with_trajectories(
          panel, model, model_predicts_change=True,
          items_to_colors=country_color_dict,
          dimensions_to_keep=(0, 1),
          n_steps_to_predict=50,
          arrow_kws={'head_length': 1, 'head_width': 1},
          predictions_plot_kws={'alpha': .9, 'linestyle': '--', 'linewidth': 5,
                                'label': 'iterated 1-year prediction'},
          empirical_plot_kws={'alpha': .3, 'linestyle': '-', 'linewidth': 3,
                              'label': 'empirical data'},
          axis_labels=None, save_fig=None):
      """Plot the model's predictions of the trajectories starting from the
      initial conditions of the items, and plot the predictions into the future
      starting from the last state of the items.
      """
      fig, ax = plt.subplots(nrows=2, figsize=(15, 15))
      items_to_trajectories = select_trajectories(panel, items_to_colors)
  
      labels_xyz_axes = compute_axes_labels(axis_labels, dimensions_to_keep)
  
      # Plot actual trajectories on both axes
      for axis in ax:
          arrow_kws_actual_trajectory = arrow_kws.copy()
          if 'alpha' in empirical_plot_kws:
              arrow_kws_actual_trajectory['alpha'] = empirical_plot_kws['alpha']
          plot_item_trajectories(
              items_to_trajectories,
              items_to_colors=items_to_colors,
              dimensions_to_keep=dimensions_to_keep,
              ax=axis,
              text_labels=True,
              arrow_kws=arrow_kws_actual_trajectory,
              plot_kws=empirical_plot_kws,
              **labels_xyz_axes)
  
      # Plot iterated predictions from the initial year for each country
      iterated_pred = iterated_predictions(
          panel, model,
          items=items_to_colors.keys(), num_time_steps='length_trajectory',
          model_predicts_change=model_predicts_change,
          index_of_initial_condition=0)
  
      plot_item_trajectories(
          iterated_pred, items_to_colors=items_to_colors,
          dimensions_to_keep=dimensions_to_keep,
          ax=ax[0], arrow_kws=arrow_kws,
          plot_kws=predictions_plot_kws)
      title = 'iterated 1-year predictions starting from the initial condition'
      ax[0].set_title(title)
  
      # Plot iterated predictions from the last year for each country
      iterated_pred = iterated_predictions(
          panel,
          model=model,
          items=items_to_trajectories.keys(),
          num_time_steps=n_steps_to_predict,
          model_predicts_change=model_predicts_change,
          index_of_initial_condition=-1)
  
      plot_item_trajectories(
          iterated_pred,
          items_to_colors=items_to_colors,
          ax=ax[1], dimensions_to_keep=dimensions_to_keep,
          arrow_kws=arrow_kws, plot_kws=predictions_plot_kws)
      title = ('iterated 1-year predictions of the next {n_steps} years '
               'starting from the last year').format(n_steps=n_steps_to_predict)
      ax[1].set_title(title)
  
      pred_patch = mpl.patches.Patch(color='k', **predictions_plot_kws)
      data_patch = mpl.patches.Patch(color='k', **empirical_plot_kws)
  
      fig.legend(
          handles=[pred_patch, data_patch],
          labels=['iterated 1-year prediction', 'empirical data'],
          loc='upper center', ncol=2, borderaxespad=0., handlelength=6)
      fig.tight_layout()
  
      maybe_save_fig(fig, save_fig)
      return fig
  
  
  def plot_iterated_predictions_for_every_item(
          panel, model, model_predicts_change=True, n_columns=6,
          figsize=(16, 40), size_of_each_plot=None,
          prediction_color=Dark2_3.mpl_colors[0],
          empirical_color=Dark2_3.mpl_colors[1],
          dimensions_to_keep=(0, 1),
          save_fig=None):
      """Plot the iterated prediction and empirical trajectory for every item.
  
      The prediction and empirical trajectories are plotted on separate axes, one
      axis for each time.
      """
      items_to_empirical_trajectories = select_trajectories(panel)
  
      iterated_pred = iterated_predictions(
          panel, model, items_to_empirical_trajectories.keys(),
          num_time_steps='length_trajectory',
          model_predicts_change=model_predicts_change,
          index_of_initial_condition=0)
  
      mse_of_final_predictions = sorted([
          (item, mean_squared_error(items_to_empirical_trajectories[item][-1],
                                    iterated_pred[item][-1]))
          for item in items_to_empirical_trajectories],
          key=lambda item_err: item_err[1])
  
      n_rows = int(np.ceil(len(items_to_empirical_trajectories) / n_columns))
      if size_of_each_plot:
          figsize = (n_columns * size_of_each_plot[0],
                     n_rows * size_of_each_plot[1])
      else:
          figsize = figsize
      fig, axes = plt.subplots(ncols=n_columns, nrows=n_rows,
                               sharex=True, sharey=True, figsize=figsize)
      axes_flat = axes.flatten()
  
      predictions_plot_kws = {
          'alpha': 1, 'linestyle': '-', 'linewidth': 2,
          'color': prediction_color}
      empirical_plot_kws = {
          'alpha': .4, 'linestyle': '-', 'linewidth': 2,
          'color': empirical_color}
  
      pred_patch = mpl.patches.Patch(**predictions_plot_kws)
      data_patch = mpl.patches.Patch(**empirical_plot_kws)
  
      for i, (item, mse_final_pred) in enumerate(mse_of_final_predictions):
          plot_item_trajectories(
              {item: iterated_pred[item]}, ax=axes_flat[i],
              plot_kws=predictions_plot_kws, arrow_start_index=-2,
              dimensions_to_keep=dimensions_to_keep,
              items_to_colors={item: prediction_color})
          plot_item_trajectories(
              {item: items_to_empirical_trajectories[item]}, ax=axes_flat[i],
              plot_kws=empirical_plot_kws, arrow_start_index=-2,
              items_to_colors={item: empirical_color},
              dimensions_to_keep=dimensions_to_keep,
              use_plot_alpha_for_arrowhead=False)
          axes_flat[i].text(.95, .95, item,
                            size=16, transform=axes_flat[i].transAxes,
                            horizontalalignment='right',
                            verticalalignment='top')
  
      fig.legend(
          handles=[pred_patch, data_patch],
          labels=['iterated 1-year prediction', 'empirical data'],
          loc='upper center', ncol=2, borderaxespad=0.)
      fig.tight_layout()
      maybe_save_fig(fig, save_fig)
      return fig
  
  
  def quiver_stream_grid(
          model, data, target, panel, model_name,
          dimensions_to_keep=(0, 1), axis_labels=None,
          n_points_each_dim=30,
          model_predicts_change=True, mask_velocity_to_convex_hull_of_data=True,
          quiver_empirical_kws=None,
          quiver_kws={'cmap': mpl.cm.viridis,
                      'angles': 'xy', 'scale_units': 'xy'},
          stream_kws={'density': 1, 'cmap': mpl.cm.viridis},
          stream_with_data_kws={'density': 0.5, 'cmap': mpl.cm.Greys},
          item_trajectories_kws=None,
          ellipse=None,
          save_fig=None):
      """Create quiver and stream plots of the time-series data and the model.
  
      Parameters
      ----------
      model : estimator with fit and predict methods that has been fit to the
      timeseries data
  
      data, target : array-like, shape [n_samples, n_features]
          `data` is the lagged times-series with lags >= 1. `target` is the
          time-series without any lags. To see how these lags are computed, see
          the function `split_panel_into_data_and_target_dataframes` in the
          module `panelyze.model_selection.split_data_target`.
  
      panel : pandas Panel
          The items are time-series; major_axis is time; minor_axis are features.
  
      model_name : str
          The name of the model to show in the plot label
  
      dimensions_to_keep : tuple of int's, each between 0 and n_features
          Which dimensions of the tensors of the meshgrid (i.e., which features
          of `data` and `target`) to show.
  
      axis_labels : None, or string, or list of strings of length equal to
      `len(dimensions_to_keep)`, or dict with keys being dimensions_to_keep and
      values being strings, default: None
          If None, then label the axes as 'dimension 0', 'dimension 1', etc.
          If a string such as 'component', then label the axes 'component 0',
          'component 1', etc.
          If a tuple of length equal to `len(dimensions_to_keep)`, then use those
          strings as the axes' labels.
          If a dict with keys given by `dimensions_to_keep`, then the labels
          are those values with '(dimension i)' (i=0, ...) written after each
          label.
  
      n_points_each_dim : int or tuple of int's of length equal to
      `len(dimensions_to_keep)`
          The number of grid points to use in each dimension for the quiver
          and stream plots
  
      model_predicts_change : bool, default: True
          Whether the model predicts the change in the time-series or the next
          value of the time-series.
  
      quiver_kws, stream_kws, stream_with_data_kws, item_trajectories_kws : dict
          Keyword arguments for the 2nd, 3rd, and 4th plots
  
  
      ellipse : None or dict {'position': (x, y), 'radius': r}
          If not None, then plot an ellipse in the bottom-left plot with given
          radius and at the given position. For example,
          ellipse={'position': (30, 16),
                   'radius': 1 / model.best_params_['gamma']**(-0.5))}
  
      save_fig : str or None, default: None
          If not None, then save the figure to the path specified by this string.
  
      Returns
      -------
      fig, axes : Figure and 4 axes
          A 2 by 2 grid of plots
      """
      item_trajectories_kws = convert_None_to_empty_dict_else_copy(
          item_trajectories_kws)
      quiver_empirical_kws = convert_None_to_empty_dict_else_copy(
          quiver_empirical_kws)
  
      fig, axes = plt.subplots(2, 2, figsize=(10, 6))
      mask_velocity = mask_velocity_to_convex_hull_of_data
      grid_points, velocities = make_data_for_streamplot(
          data, model, n_points_each_dim,
          model_predicts_change=model_predicts_change,
          dimensions_to_keep=dimensions_to_keep,
          mask_velocity_to_convex_hull_of_data=mask_velocity)
  
      labels_xyz_axes = compute_axes_labels(axis_labels, dimensions_to_keep)
  
      quiver_plot_empirical(data, target, dimensions_to_keep=dimensions_to_keep,
                            ax=axes[0, 0],
                            **quiver_empirical_kws, **labels_xyz_axes)
      quiver_plot(
          *itertools.chain(np.meshgrid(*grid_points, indexing='ij'), velocities),
          ax=axes[0, 1], quiver_kws=quiver_kws, **labels_xyz_axes)
  
      stream_plot(*itertools.chain(grid_points, velocities),
                  stream_kws=stream_kws, ax=axes[1, 0], **labels_xyz_axes)
      if ellipse:
          plot_ellipse(axes[1, 0], ellipse['position'], ellipse['radius'])
      stream_plot(*itertools.chain(grid_points, velocities),
                  stream_kws=stream_with_data_kws, ax=axes[1, 1],
                  **labels_xyz_axes)
      plot_item_trajectories(
          panel[panel.items.isin(country_color_dict)],
          ax=axes[1, 1], dimensions_to_keep=dimensions_to_keep,
          text_labels=True,
          **item_trajectories_kws)
      title = 'data and predictions for {}'.format(model_name)
      fig.suptitle(title, y=1.02)
      fig.tight_layout()
      maybe_save_fig(fig, save_fig)
      return fig, axes
  
  
  def make_axis_labels(axis_labels_dict, dimensions_to_keep, label='dimension'):
      """Make a dictionary of axis label names {'xlabel': 'dimension 0', ...}.
  
      Parameters
      ----------
      axis_labels_dict : None or dict
          A dictionary mapping dimension indices to strings, such as
          {0: 'component 0', 1: 'component 1'}.
          If None, then use '{label} i' where '{label}' is replaced by the
          keyword argument `label`. If any dimension in `dimensions_to_keep` is
          not a key of `axis_labels_dict`, then the default label '{label} i'
          is used.
  
      dimensions_to_keep : list of int's
          The integers of which dimensions are being selected.
  
      label : str, default: 'dimension'
          The string to use to label axes when axis_labels_dict is None.
  
      Returns
      -------
      axis_label_names_to_axis_labels : dict
          Dictionary with keys in ['xlabel', 'ylabel', 'zlabel'], and values are
          strings.
      """
      axis_label_names = ['xlabel', 'ylabel', 'zlabel']
  
      if axis_labels_dict is None:
          axis_labels_dict = {}
      axis_labels = [axis_labels_dict.get(dim, '{} {}'.format(label, dim))
                     for dim in dimensions_to_keep]
      return dict(zip(axis_label_names, axis_labels))
  
  
  def compute_axes_labels(axis_labels, dimensions_to_keep, label='dimension'):
      axis_label_names = ['xlabel', 'ylabel', 'zlabel']
  
      if isinstance(axis_labels, str):
          label = axis_labels
  
      default_axis_labels = ['{label} {dim}'.format(label=label, dim=dim)
                             for dim in dimensions_to_keep]
  
      if axis_labels is None or isinstance(axis_labels, str):
          axis_labels = default_axis_labels
      elif isinstance(axis_labels, dict):
          axis_labels = [axis_labels[dim] for dim in dimensions_to_keep]
          axis_labels = ['{} ({})'.format(l, d)
                         for l, d in zip(axis_labels, default_axis_labels)]
      elif len(axis_labels) == len(dimensions_to_keep):
          pass
      else:
          msg = ('Could not compute axes labels from {al} with '
                 '`dimensions_to_keep` = {dtk}')
          raise ValueError(msg.format(al=axis_labels, dtk=dimensions_to_keep))
  
      return dict(zip(axis_label_names, axis_labels))
  
  
  def density_list_plot_2d(
          x, y, z, convex_hull_only=True, error_label='error',
          n_points_x=100, n_points_y=100,
          xlabel='score on the first principal component',
          ylabel='score on the second principal component',
          log_color=False, subplots_kws=None, ax=None, save_fig=None):
      """Return a density plot of z as a color, interpolated with an RBF kernel.
  
      Parameters
      ----------
      x, y, z : arrays of the same length
  
      convex_hull_only : bool, default: True
          Whether to only plot values in the convex hull of the `x` and `y` data
  
      error_label : str, default: 'error'
          The label on the colorbar
  
      n_points_x, n_points_y : int, default: 100
          The number of bins to use in the x and y dimensions
  
      xlabel, ylabel : str
          Axis labels
  
      log_color : bool, default: True
          Whether to logarithmically transform the color scale
  
      ax : matplotlib axis, optional, default: None
          The axis on which to draw the plot. If None, then an axis is created.
  
      save_fig : str or None, default: None
          If not None, then save the figure to the path specified by this string.
  
      subplots_kws : keyword arguments to pass to plt.subplots, default: None
      """
      # Set up a regular grid of interpolation points
      xi = np.linspace(x.min(), x.max(), n_points_x)
      yi = np.linspace(y.min(), y.max(), n_points_y)
      xgrid, ygrid = np.meshgrid(xi, yi)
  
      # Interpolate
      rbf = scipy.interpolate.Rbf(x, y, z, function='linear')
      zi = rbf(xgrid, ygrid)
  
      if convex_hull_only:
          zi = mask_arrays_with_convex_hull(
              [zi], [xi, yi],
              scipy.spatial.ConvexHull(
                  np.vstack((x, y)).T))[0]
  
      subplots_kws = convert_None_to_empty_dict_else_copy(subplots_kws)
      fig, ax = create_fig_ax(ax, **subplots_kws)
      axim = ax.imshow(
          zi, vmin=z.min(), vmax=z.max(), origin='lower',
          extent=[x.min(), x.max(), y.min(), y.max()],
          norm=(mpl.colors.LogNorm(vmin=zi.min(), vmax=zi.max()) if log_color
                else None))
  
      ax.set_xlabel(xlabel)
      ax.set_ylabel(ylabel)
  
      cax, kwargs = mpl.colorbar.make_axes(
          [ax], location='bottom', fraction=0.05, pad=.1)
      cax.text(-.15, 0.1, error_label, ha='center', va='center')
      fig.colorbar(axim, cax=cax, orientation='horizontal', extend='max')
  
      maybe_save_fig(fig, save_fig)
      return fig, ax
  
apis: sklearn.utils.check_array sklearn.metrics.mean_squared_error sklearn.kernel_ridge.KernelRidge