text: # Different augmentation functions
  import nengo
  from nengo.processes import WhiteSignal, WhiteNoise
  from dataman import d3_scale
  from constants import *
  import numpy as np
  
  import json
  import itertools
  
  import matplotlib.pyplot as plt
  import ipdb
  
  
  # TODO: get rid of the amount argument
  def aug(dataset, desc, amount, func, kwargs=None, save_dir="../datasets"):
      """iterate through each class, augment it and save the result.
  
      The `amount` argument is to note the number of examples"""
      kwargs = kwargs or {}
      new_data = []
  
      for _ in range(desc["n_classes"]):
          new_data.append(
              np.zeros((amount, desc["dims"], int(desc["t_len"]/dt)))
          )
  
      for c_i, s_i, d_i in itertools.product(xrange(len(dataset)), xrange(amount), xrange(desc["dims"])):
          new_data[c_i][s_i, d_i] = func(dataset[c_i][s_i][d_i], desc["t_len"], **kwargs)
  
      if save_dir is not None:
          filename = "%s/dataset_%scls_%s_%s_%s_%s_aug_%s" % (
              save_dir,
              desc["class_type"],
              desc["t_len"],
              desc["dims"],
              desc["n_classes"],
              desc["SEED"],
              func.__name__
          )
  
          # Maintain a list of all the augmented files using json
          with open("%s/aug_log.json" % save_dir, "r") as f_log:
              file_list = json.loads(f_log.read())
          file_list[filename] = {"desc": desc, "func": func.__name__, "kwargs": kwargs}
          with open("%s/aug_log.json" % save_dir, "w") as f_log:
              f_log.write(json.dumps(file_list))
  
          # TODO: give option to append to an existing file if the filename already exists
          # especially with the same kwargs
          np.savez(filename, class_sig_list=new_data, class_desc=desc)
  
      return new_data
  
  
  def add_rand_noise(dataset, t_len, freq=10, scale=0.2, sig=True):
      """additive noise"""
      if sig:
          noise = WhiteSignal(t_len, freq).run(t_len)[:, 0] * scale
      else:
          noise = WhiteNoise().run(t_len)[:, 0] * scale
      return dataset + noise
  
  
  def conv_rand_noise(dataset, t_len, freq=500, scale=0.001):
      """convolve with noise, which heavily distorts the signal almost beyond all recognition"""
      noise = WhiteSignal(t_len, freq).run(t_len)[:, 0] * scale
      return d3_scale(np.convolve(dataset, noise, mode="same"))
  
  
  def shot_noise(dataset, t_len, shots=3, width=2):
      """randomly add impulses in either direction"""
      t_len = t_len/dt
      assert shots*width < t_len
      assert t_len % width == 0
      shot_ind = np.random.choice(int(t_len/width), replace=False, size=shots)
      shot_vals = np.random.choice([-1, 1], size=(shots, 1))
  
      # reshape into widths, index to add then flatten again
      ret_val = np.copy(dataset)
      ret_val.reshape((-1, width))[shot_ind, :] += shot_vals
      # no need to rescale
      return ret_val.flatten()
  
  
  def low_filt(dataset, t_len, tau=0.01):
      return nengo.Lowpass(tau).filtfilt(dataset)
  
  
  def offset(dataset, t_len, scale=0.1):
      """shift and normalize"""
      if scale > 0.1:
          out_range = (np.min(dataset), np.max(dataset+scale))
      else:
          out_range = (np.min(dataset+scale), np.max(dataset))
  
      return d3_scale(dataset + scale, out_range=out_range)
  
  
  def make_more(dataset):
      # won't work for '_spec' signals
      raise NotImplementedError("Nope")
  
  
  def lag(dataset, t_len, lags=3, width=10):
      """add a random temporal lag"""
  
      raise NotImplementedError("Nope")
      t_len = t_len/dt
      assert lags*width < t_len
      assert t_len % width == 0
      lag_ind = np.random.choice(int(t_len/width), replace=False, size=lags)
      ret_val = np.copy(dataset)
  
      # reshape into widths, index to add then flatten again
      lag_vals = ret_val.reshape((-1, width))[lag_ind, 0]
      ret_val.reshape((-1, width))[lag_ind, :] = (np.ones((width, lags)) * lag_vals).T
      # no need to rescale
      return ret_val.flatten()
  
  
  def pre_arggen_repeat(dat, repeats=3):
      """Repeat the dataset
  
      Expects data from a created dataset unformatted for specific network input"""
  
      final_dat = np.tile(dat, (1, repeats, 1, 1))
      return final_dat
  
  
  def post_arggen_repeat(dat, cor, repeats=3):
      """Repeat the dataset
  
      Expects data already formatted for Nengo input"""
  
      final_dat = np.tile(dat, (repeats, 1, 1))
      final_cor = np.tile(cor, (repeats, 1))
  
      return final_dat, final_cor
  
  
  def dat_shuffle(dat, cor, rng=np.random.RandomState(SEED)):
      idx = np.arange(cor.shape[0])
      rng.shuffle(idx)
      return dat[idx], cor[idx]
  
  
  def dat_repshuf(dat, cor, reps=3, rng=np.random.RandomState(SEED)):
      r_dat, r_cor = post_arggen_repeat(dat, cor, reps)
      return dat_shuffle(r_dat, r_cor, rng)
  