<commit_before>from dapper import *

#from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import juggle_axes

from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

from matplotlib import colors
from matplotlib.ticker import MaxNLocator


def setup_wrapping(M,periodic=True):
  """
  Make periodic indices and a corresponding function
  (that works for ensemble input).
  """

  if periodic:
    ii = np.hstack([-0.5, arange(M), M-0.5])
    def wrap(E):
      midpoint = (E[[0],...] + E[[-1],...])/2
      return ccat(midpoint,E,midpoint)

  else:
    ii = arange(M)
    wrap = lambda x: x

  return ii, wrap
  
def adjust_position(ax,adjust_extent=False,**kwargs):
  """
  Adjust values (add) to get_position().
  kwarg must be one of 'x0','y0','width','height'.
  """
  # Load get_position into d
  pos = ax.get_position()
  d   = OrderedDict()
  for key in ['x0','y0','width','height']:
    d[key] = getattr(pos,key)
  # Make adjustments
  for key,item in kwargs.items():
    d[key] += item
    if adjust_extent:
      if key=='x0': d['width']  -= item
      if key=='y0': d['height'] -= item
  # Set
  ax.set_position(d.values())

def xtrema(xx,axis=None):
  a = xx.min(axis)
  b = xx.max(axis)
  return a, b

def stretch(a,b,factor=1,int=False):
  """
  Stretch distance a-b by factor.
  Return a,b.
  If int: floor(a) and ceil(b)
  """
  c = (a+b)/2
  a = c + factor*(a-c) 
  b = c + factor*(b-c) 
  if int:
    a = floor(a)
    b = ceil(b)
  return a, b


def set_ilim(ax,i,Min=None,Max=None):
  """Set bounds on axis i.""" 
  if i is 0: ax.set_xlim(Min,Max)
  if i is 1: ax.set_ylim(Min,Max)
  if i is 2: ax.set_zlim(Min,Max)

# Examples:
# K_lag = estimate_good_plot_length(stats.xx,chrono,mult = 80)
def estimate_good_plot_length(xx,chrono=None,mult=100):
  """
  Estimate good length for plotting stuff
  from the time scale of the system.
  Provide sensible fall-backs (better if chrono is supplied).
  """
  if xx.ndim == 2:
    # If mult-dim, then average over dims (by ravel)....
    # But for inhomogeneous variables, it is important
    # to subtract the mean first!
    xx = xx - mean(xx,axis=0)
    xx = xx.ravel(order='F')

  try:
    K = mult * estimate_corr_length(xx)
  except ValueError:
    K = 0

  if chrono != None:
    t = chrono
    K = int(min(max(K, t.dkObs), t.K))
    T = round2sigfig(t.tt[K],2) # Could return T; T>tt[-1]
    K = find_1st_ind(t.tt >= T)
    if K: return K
    else: return t.K
  else:
    K = int(min(max(K, 1), len(xx)))
    T = round2sigfig(K,2)
    return K


def plot_pause(interval):
  """Similar to plt.pause()"""

  # plt.pause(0) just seems to freeze execution.
  if interval==0:
    return

  try:
    # Implement plt.pause() that doesn't focus window, c.f.
    # github.com/matplotlib/matplotlib/issues/11131, so.com/q/45729092.
    # Only necessary for some platforms (e.g. Windows) and mpl versions.
    # Even then, mere figure creation may steal the focus.
    # This was done deliberately github.com/matplotlib/matplotlib/pull/6384#issue-69259165.
    # See issue: github.com/matplotlib/matplotlib/issues/8246#issuecomment-505460935
    from matplotlib import _pylab_helpers
    def _plot_pause(interval,  focus_figure=True):
        canvas = plt.gcf().canvas
        manager = canvas.manager
        if manager is not None:
            if canvas.figure.stale:
                canvas.draw_idle()
            if focus_figure:
                plt.show(block=False)
            # if not is_notebook: # also see below
            if True:
                canvas.start_event_loop(interval)
        else:
            time.sleep(interval)
    _plot_pause(interval, focus_figure=False)

  except:
    # Jupyter notebook support (SO.com/q/34486642)
    # Note: no longer needed with the above _plot_pause()?
    plt.gcf().canvas.draw()
    time.sleep(0.1)


def plot_hovmoller(xx,chrono=None,**kwargs):
  """
  Plot Hovmöller diagram.
  """
  fig, ax = freshfig(26,figsize=(4,3.5),loc='331-22')

  if chrono!=None:
    mask = chrono.tt <= chrono.Tplot*2
    kk   = chrono.kk[mask]
    tt   = chrono.tt[mask]
    ax.set_ylabel('Time (t)')
  else:
    K    = estimate_good_plot_length(xx,mult=20)
    kk   = arange(K)
    tt   = kk
    ax.set_ylabel('Time indices (k)')

  plt.contourf(arange(xx.shape[1]),tt,xx[kk],25)
  plt.colorbar()
  ax.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  ax.set_title("Hovmoller diagram (of 'Truth')")
  ax.set_xlabel('Dimension index (i)')

  plot_pause(0.1)
  plt.tight_layout()


def add_endpoint_xtick(ax):
  """Useful when xlim(right) is e.g. 39 (instead of 40)."""
  xF = ax.get_xlim()[1]
  ticks = ax.get_xticks()
  if ticks[-1] > xF:
    ticks = ticks[:-1]
  ticks = np.append(ticks, xF)
  ax.set_xticks(ticks)


def integer_hist(E,N,centrd=False,weights=None,**kwargs):
  """Histogram for integers."""
  ax = plt.gca()
  rnge = (-0.5,N+0.5) if centrd else (0,N+1)
  ax.hist(E,bins=N+1,range=rnge,density=True,weights=weights,**kwargs)
  ax.set_xlim(rnge)


def not_available_text(ax,txt=None,fs=20):
  if txt is None: txt = '[Not available]'
  else:           txt = '[' + txt + ']'
  ax.text(0.5,0.5,txt,
      fontsize=fs,
      transform=ax.transAxes,
      va='center',ha='center',
      wrap=True)

def plot_err_components(stats):
  """Plot components of the error.

  .. note::
    it was chosen to plot(ii, mean_in_time(abs(err_i))),
    and thus the corresponding spread measure is MAD.
    If one chose instead: plot(ii, std_in_time(err_i)),
    then the corresponding measure of spread would have been std.
    This choice was made in part because (wrt. subplot 2)
    the singular values (svals) correspond to rotated MADs,
    and because rms(umisf) seems to convoluted for interpretation.
  """
  fig, (ax0,ax1,ax2) = freshfig(25,figsize=(6,6),loc='1313',nrows=3)

  chrono = stats.HMM.t
  Nx     = stats.xx.shape[1]

  err   = mean( abs(stats.err  .a) ,0)
  sprd  = mean(     stats.mad  .a  ,0)
  umsft = mean( abs(stats.umisf.a) ,0)
  usprd = mean(     stats.svals.a  ,0)

  ax0.plot(          arange(Nx),               err,'k',lw=2, label='Error')
  if Nx<10**3:
    ax0.fill_between(arange(Nx),[0]*len(sprd),sprd,alpha=0.7,label='Spread')
  else:
    ax0.plot(        arange(Nx),              sprd,alpha=0.7,label='Spread')
  #ax0.set_yscale('log')
  ax0.set_title('Element-wise error comparison')
  ax0.set_xlabel('Dimension index (i)')
  ax0.set_ylabel('Time-average (_a) magnitude')
  ax0.set_xlim(0,Nx-1)
  ax0.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  ax0.legend(loc='upper right')

  ax1.set_xlim(0,Nx-1)
  ax1.set_xlabel('Principal component index')
  ax1.set_ylabel('Time-average (_a) magnitude')
  ax1.set_title('Spectral error comparison')
  has_been_computed = np.any(np.isfinite(umsft))
  if has_been_computed:
    L = len(umsft)
    ax1.plot(        arange(L),      umsft,'k',lw=2, label='Error')
    ax1.fill_between(arange(L),[0]*L,usprd,alpha=0.7,label='Spread')
    ax1.set_yscale('log')
    ax1.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  else:
    not_available_text(ax1)

  rmse = stats.rmse.a[chrono.maskObs_BI]
  ax2.hist(rmse,bins=30,density=False)
  ax2.set_ylabel('Num. of occurence (_a)')
  ax2.set_xlabel('RMSE')
  ax2.set_title('Histogram of RMSE values')
  ax2.set_xlim(left=0)

  plot_pause(0.1)
  plt.tight_layout()


def plot_rank_histogram(stats):
  chrono = stats.HMM.t

  has_been_computed = \
      hasattr(stats,'rh') and \
      not all(stats.rh.a[-1]==array(np.nan).astype(int))

  fig, ax = freshfig(24, (6,3), loc="3313")
  ax.set_title('(Mean of marginal) rank histogram (_a)')
  ax.set_ylabel('Freq. of occurence\n (of truth in interval n)')
  ax.set_xlabel('ensemble member index (n)')

  if has_been_computed:
    ranks = stats.rh.a[chrono.maskObs_BI]
    Nx    = ranks.shape[1]
    N     = stats.config.N
    if not hasattr(stats,'w'):
      # Ensemble rank histogram
      integer_hist(ranks.ravel(),N)
    else:
      # Experimental: weighted rank histogram.
      # Weight ranks by inverse of particle weight. Why? Coz, with correct
      # importance weights, the "expected value" histogram is then flat.
      # Potential improvement: interpolate weights between particles.
      w  = stats.w.a[chrono.maskObs_BI]
      K  = len(w)
      w  = np.hstack([w, ones((K,1))/N]) # define weights for rank N+1
      w  = array([ w[arange(K),ranks[arange(K),i]] for i in range(Nx)])
      w  = w.T.ravel()
      w  = np.maximum(w, 1/N/100) # Artificial cap. Reduces variance, but introduces bias.
      w  = 1/w
      integer_hist(ranks.ravel(),N,weights=w)
  else:
    not_available_text(ax)
  
  plot_pause(0.1)
  plt.tight_layout()


def adjustable_box_or_forced():
  "For set_aspect(), adjustable='box-forced' replaced by 'box' since mpl 2.2.0."
  from pkg_resources import parse_version as pv
  return 'box-forced' if pv(mpl.__version__) < pv("2.2.0") else 'box'


def freshfig(num,figsize=None,*args,**kwargs):
  """Create/clear figure.
  
  Similar to::

    fig, ax = suplots(*args,**kwargs)

  With the modification that:

  - If the figure does not exist: create it.
    This allows for figure sizing -- even with mpl backend MacOS.
    Also place figure.<commit_msg>Set None as default num in freshfig<commit_after>from dapper import *

#from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import juggle_axes

from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

from matplotlib import colors
from matplotlib.ticker import MaxNLocator


def setup_wrapping(M,periodic=True):
  """
  Make periodic indices and a corresponding function
  (that works for ensemble input).
  """

  if periodic:
    ii = np.hstack([-0.5, arange(M), M-0.5])
    def wrap(E):
      midpoint = (E[[0],...] + E[[-1],...])/2
      return ccat(midpoint,E,midpoint)

  else:
    ii = arange(M)
    wrap = lambda x: x

  return ii, wrap
  
def adjust_position(ax,adjust_extent=False,**kwargs):
  """
  Adjust values (add) to get_position().
  kwarg must be one of 'x0','y0','width','height'.
  """
  # Load get_position into d
  pos = ax.get_position()
  d   = OrderedDict()
  for key in ['x0','y0','width','height']:
    d[key] = getattr(pos,key)
  # Make adjustments
  for key,item in kwargs.items():
    d[key] += item
    if adjust_extent:
      if key=='x0': d['width']  -= item
      if key=='y0': d['height'] -= item
  # Set
  ax.set_position(d.values())

def xtrema(xx,axis=None):
  a = xx.min(axis)
  b = xx.max(axis)
  return a, b

def stretch(a,b,factor=1,int=False):
  """
  Stretch distance a-b by factor.
  Return a,b.
  If int: floor(a) and ceil(b)
  """
  c = (a+b)/2
  a = c + factor*(a-c) 
  b = c + factor*(b-c) 
  if int:
    a = floor(a)
    b = ceil(b)
  return a, b


def set_ilim(ax,i,Min=None,Max=None):
  """Set bounds on axis i.""" 
  if i is 0: ax.set_xlim(Min,Max)
  if i is 1: ax.set_ylim(Min,Max)
  if i is 2: ax.set_zlim(Min,Max)

# Examples:
# K_lag = estimate_good_plot_length(stats.xx,chrono,mult = 80)
def estimate_good_plot_length(xx,chrono=None,mult=100):
  """
  Estimate good length for plotting stuff
  from the time scale of the system.
  Provide sensible fall-backs (better if chrono is supplied).
  """
  if xx.ndim == 2:
    # If mult-dim, then average over dims (by ravel)....
    # But for inhomogeneous variables, it is important
    # to subtract the mean first!
    xx = xx - mean(xx,axis=0)
    xx = xx.ravel(order='F')

  try:
    K = mult * estimate_corr_length(xx)
  except ValueError:
    K = 0

  if chrono != None:
    t = chrono
    K = int(min(max(K, t.dkObs), t.K))
    T = round2sigfig(t.tt[K],2) # Could return T; T>tt[-1]
    K = find_1st_ind(t.tt >= T)
    if K: return K
    else: return t.K
  else:
    K = int(min(max(K, 1), len(xx)))
    T = round2sigfig(K,2)
    return K


def plot_pause(interval):
  """Similar to plt.pause()"""

  # plt.pause(0) just seems to freeze execution.
  if interval==0:
    return

  try:
    # Implement plt.pause() that doesn't focus window, c.f.
    # github.com/matplotlib/matplotlib/issues/11131, so.com/q/45729092.
    # Only necessary for some platforms (e.g. Windows) and mpl versions.
    # Even then, mere figure creation may steal the focus.
    # This was done deliberately github.com/matplotlib/matplotlib/pull/6384#issue-69259165.
    # See issue: github.com/matplotlib/matplotlib/issues/8246#issuecomment-505460935
    from matplotlib import _pylab_helpers
    def _plot_pause(interval,  focus_figure=True):
        canvas = plt.gcf().canvas
        manager = canvas.manager
        if manager is not None:
            if canvas.figure.stale:
                canvas.draw_idle()
            if focus_figure:
                plt.show(block=False)
            # if not is_notebook: # also see below
            if True:
                canvas.start_event_loop(interval)
        else:
            time.sleep(interval)
    _plot_pause(interval, focus_figure=False)

  except:
    # Jupyter notebook support (SO.com/q/34486642)
    # Note: no longer needed with the above _plot_pause()?
    plt.gcf().canvas.draw()
    time.sleep(0.1)


def plot_hovmoller(xx,chrono=None,**kwargs):
  """
  Plot Hovmöller diagram.
  """
  fig, ax = freshfig(26,figsize=(4,3.5),loc='331-22')

  if chrono!=None:
    mask = chrono.tt <= chrono.Tplot*2
    kk   = chrono.kk[mask]
    tt   = chrono.tt[mask]
    ax.set_ylabel('Time (t)')
  else:
    K    = estimate_good_plot_length(xx,mult=20)
    kk   = arange(K)
    tt   = kk
    ax.set_ylabel('Time indices (k)')

  plt.contourf(arange(xx.shape[1]),tt,xx[kk],25)
  plt.colorbar()
  ax.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  ax.set_title("Hovmoller diagram (of 'Truth')")
  ax.set_xlabel('Dimension index (i)')

  plot_pause(0.1)
  plt.tight_layout()


def add_endpoint_xtick(ax):
  """Useful when xlim(right) is e.g. 39 (instead of 40)."""
  xF = ax.get_xlim()[1]
  ticks = ax.get_xticks()
  if ticks[-1] > xF:
    ticks = ticks[:-1]
  ticks = np.append(ticks, xF)
  ax.set_xticks(ticks)


def integer_hist(E,N,centrd=False,weights=None,**kwargs):
  """Histogram for integers."""
  ax = plt.gca()
  rnge = (-0.5,N+0.5) if centrd else (0,N+1)
  ax.hist(E,bins=N+1,range=rnge,density=True,weights=weights,**kwargs)
  ax.set_xlim(rnge)


def not_available_text(ax,txt=None,fs=20):
  if txt is None: txt = '[Not available]'
  else:           txt = '[' + txt + ']'
  ax.text(0.5,0.5,txt,
      fontsize=fs,
      transform=ax.transAxes,
      va='center',ha='center',
      wrap=True)

def plot_err_components(stats):
  """Plot components of the error.

  .. note::
    it was chosen to plot(ii, mean_in_time(abs(err_i))),
    and thus the corresponding spread measure is MAD.
    If one chose instead: plot(ii, std_in_time(err_i)),
    then the corresponding measure of spread would have been std.
    This choice was made in part because (wrt. subplot 2)
    the singular values (svals) correspond to rotated MADs,
    and because rms(umisf) seems to convoluted for interpretation.
  """
  fig, (ax0,ax1,ax2) = freshfig(25,figsize=(6,6),loc='1313',nrows=3)

  chrono = stats.HMM.t
  Nx     = stats.xx.shape[1]

  err   = mean( abs(stats.err  .a) ,0)
  sprd  = mean(     stats.mad  .a  ,0)
  umsft = mean( abs(stats.umisf.a) ,0)
  usprd = mean(     stats.svals.a  ,0)

  ax0.plot(          arange(Nx),               err,'k',lw=2, label='Error')
  if Nx<10**3:
    ax0.fill_between(arange(Nx),[0]*len(sprd),sprd,alpha=0.7,label='Spread')
  else:
    ax0.plot(        arange(Nx),              sprd,alpha=0.7,label='Spread')
  #ax0.set_yscale('log')
  ax0.set_title('Element-wise error comparison')
  ax0.set_xlabel('Dimension index (i)')
  ax0.set_ylabel('Time-average (_a) magnitude')
  ax0.set_xlim(0,Nx-1)
  ax0.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  ax0.legend(loc='upper right')

  ax1.set_xlim(0,Nx-1)
  ax1.set_xlabel('Principal component index')
  ax1.set_ylabel('Time-average (_a) magnitude')
  ax1.set_title('Spectral error comparison')
  has_been_computed = np.any(np.isfinite(umsft))
  if has_been_computed:
    L = len(umsft)
    ax1.plot(        arange(L),      umsft,'k',lw=2, label='Error')
    ax1.fill_between(arange(L),[0]*L,usprd,alpha=0.7,label='Spread')
    ax1.set_yscale('log')
    ax1.get_xaxis().set_major_locator(MaxNLocator(integer=True))
  else:
    not_available_text(ax1)

  rmse = stats.rmse.a[chrono.maskObs_BI]
  ax2.hist(rmse,bins=30,density=False)
  ax2.set_ylabel('Num. of occurence (_a)')
  ax2.set_xlabel('RMSE')
  ax2.set_title('Histogram of RMSE values')
  ax2.set_xlim(left=0)

  plot_pause(0.1)
  plt.tight_layout()


def plot_rank_histogram(stats):
  chrono = stats.HMM.t

  has_been_computed = \
      hasattr(stats,'rh') and \
      not all(stats.rh.a[-1]==array(np.nan).astype(int))

  fig, ax = freshfig(24, (6,3), loc="3313")
  ax.set_title('(Mean of marginal) rank histogram (_a)')
  ax.set_ylabel('Freq. of occurence\n (of truth in interval n)')
  ax.set_xlabel('ensemble member index (n)')

  if has_been_computed:
    ranks = stats.rh.a[chrono.maskObs_BI]
    Nx    = ranks.shape[1]
    N     = stats.config.N
    if not hasattr(stats,'w'):
      # Ensemble rank histogram
      integer_hist(ranks.ravel(),N)
    else:
      # Experimental: weighted rank histogram.
      # Weight ranks by inverse of particle weight. Why? Coz, with correct
      # importance weights, the "expected value" histogram is then flat.
      # Potential improvement: interpolate weights between particles.
      w  = stats.w.a[chrono.maskObs_BI]
      K  = len(w)
      w  = np.hstack([w, ones((K,1))/N]) # define weights for rank N+1
      w  = array([ w[arange(K),ranks[arange(K),i]] for i in range(Nx)])
      w  = w.T.ravel()
      w  = np.maximum(w, 1/N/100) # Artificial cap. Reduces variance, but introduces bias.
      w  = 1/w
      integer_hist(ranks.ravel(),N,weights=w)
  else:
    not_available_text(ax)
  
  plot_pause(0.1)
  plt.tight_layout()


def adjustable_box_or_forced():
  "For set_aspect(), adjustable='box-forced' replaced by 'box' since mpl 2.2.0."
  from pkg_resources import parse_version as pv
  return 'box-forced' if pv(mpl.__version__) < pv("2.2.0") else 'box'


def freshfig(num=None,figsize=None,*args,**kwargs):
  """Create/clear figure.
  
  Similar to::

    fig, ax = suplots(*args,**kwargs)

  With the modification that:

  - If the figure does not exist: create it.
    This allows for figure sizing -- even with mpl backend MacOS.
    Also place figure.
{'fraction_non_alphanumeric': 0.1447361328230284, 'fraction_numerical': 0.023287154331302896, 'mean_word_length': 3.2192403911758016, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5352644', 'n_tokens_mistral': 7441, 'n_tokens_neox': 7066, 'n_words': 1597}