import numpy as np
from jug import TaskGenerator
import ezplot


__all__ = ['plot_stack']


@TaskGenerator
def plot_stack(stack_results, problems=None, policies=None, name=''):
    problems = problems if problems is not None else stack_results.key()
    nfigs = len(problems)

    fig = ezplot.figure(figsize=(5*nfigs, 4))

    for i, expt in enumerate(problems):
        results = stack_results[expt]
        policies = policies if policies is not None else results.keys()

        ax = fig.add_subplot(1, nfigs, i+1)

        for policy in policies:
            xbest, ybest = zip(*results[policy])
            iters = np.arange(np.shape(ybest)[1])

            mu = np.mean(ybest, axis=0)
            std = np.std(ybest, axis=0) / np.sqrt(len(ybest))
            ax.plot_banded(iters, mu, std, label=policy)

        ax.set_title(expt, fontsize=16)

    ax.legend(loc=0, fontsize=16)
    ezplot.plt.savefig(name)

    return fig


import matplotlib
matplotlib.use('Agg')
import numpy as np
from jug import TaskGenerator
import ezplot


__all__ = ['plot_stack']


@TaskGenerator
def plot_stack(stack_results, problems=None, policies=None, name=''):
    problems = problems if problems is not None else stack_results.key()
    nfigs = len(problems)

    fig = ezplot.figure(figsize=(5*nfigs, 4))

    for i, expt in enumerate(problems):
        results = stack_results[expt]
        policies = policies if policies is not None else results.keys()

        ax = fig.add_subplot(1, nfigs, i+1)

        for policy in policies:
            xbest, ybest = zip(*results[policy])
            iters = np.arange(np.shape(ybest)[1])

            mu = np.mean(ybest, axis=0)
            std = np.std(ybest, axis=0) / np.sqrt(len(ybest))
            ax.plot_banded(iters, mu, std, label=policy)

        ax.set_title(expt, fontsize=16)

    ax.legend(loc=0, fontsize=16)
    ezplot.plt.savefig(name)

    return fig


Fix plotter to use 'Agg'.
Fix plotter to use 'Agg'.

+ import matplotlib
+ matplotlib.use('Agg')
  import numpy as np
  from jug import TaskGenerator
  import ezplot
  
  
  __all__ = ['plot_stack']
  
  
  @TaskGenerator
  def plot_stack(stack_results, problems=None, policies=None, name=''):
      problems = problems if problems is not None else stack_results.key()
      nfigs = len(problems)
  
      fig = ezplot.figure(figsize=(5*nfigs, 4))
  
      for i, expt in enumerate(problems):
          results = stack_results[expt]
          policies = policies if policies is not None else results.keys()
  
          ax = fig.add_subplot(1, nfigs, i+1)
  
          for policy in policies:
              xbest, ybest = zip(*results[policy])
              iters = np.arange(np.shape(ybest)[1])
  
              mu = np.mean(ybest, axis=0)
              std = np.std(ybest, axis=0) / np.sqrt(len(ybest))
              ax.plot_banded(iters, mu, std, label=policy)
  
          ax.set_title(expt, fontsize=16)
  
      ax.legend(loc=0, fontsize=16)
      ezplot.plt.savefig(name)
  
      return fig
  
  
Fix plotter to use 'Agg'.
## Code Before:
import numpy as np
from jug import TaskGenerator
import ezplot


__all__ = ['plot_stack']


@TaskGenerator
def plot_stack(stack_results, problems=None, policies=None, name=''):
    problems = problems if problems is not None else stack_results.key()
    nfigs = len(problems)

    fig = ezplot.figure(figsize=(5*nfigs, 4))

    for i, expt in enumerate(problems):
        results = stack_results[expt]
        policies = policies if policies is not None else results.keys()

        ax = fig.add_subplot(1, nfigs, i+1)

        for policy in policies:
            xbest, ybest = zip(*results[policy])
            iters = np.arange(np.shape(ybest)[1])

            mu = np.mean(ybest, axis=0)
            std = np.std(ybest, axis=0) / np.sqrt(len(ybest))
            ax.plot_banded(iters, mu, std, label=policy)

        ax.set_title(expt, fontsize=16)

    ax.legend(loc=0, fontsize=16)
    ezplot.plt.savefig(name)

    return fig


## Instruction:
Fix plotter to use 'Agg'.
## Code After:
import matplotlib
matplotlib.use('Agg')
import numpy as np
from jug import TaskGenerator
import ezplot


__all__ = ['plot_stack']


@TaskGenerator
def plot_stack(stack_results, problems=None, policies=None, name=''):
    problems = problems if problems is not None else stack_results.key()
    nfigs = len(problems)

    fig = ezplot.figure(figsize=(5*nfigs, 4))

    for i, expt in enumerate(problems):
        results = stack_results[expt]
        policies = policies if policies is not None else results.keys()

        ax = fig.add_subplot(1, nfigs, i+1)

        for policy in policies:
            xbest, ybest = zip(*results[policy])
            iters = np.arange(np.shape(ybest)[1])

            mu = np.mean(ybest, axis=0)
            std = np.std(ybest, axis=0) / np.sqrt(len(ybest))
            ax.plot_banded(iters, mu, std, label=policy)

        ax.set_title(expt, fontsize=16)

    ax.legend(loc=0, fontsize=16)
    ezplot.plt.savefig(name)

    return fig

