content: #!/usr/bin/python
  """code to help run the image-computable version of the model
  
  we're using this primarily to check the effect of vignetting, but this does make our project
  image-computable (though it's a linear model and so will fail in some trivial cases)
  
  """
  import itertools
  import argparse
  import numpy as np
  import pandas as pd
  import pyrtools as pt
  from scipy import interpolate
  
  
  def upsample(signal, target_shape):
      """upsample a signal to target_shape
  
      this uses scipy's interpolate.interp2d (and so will end up with a smoothed signal)
      """
      x = np.linspace(-(signal.shape[0]-1)/2, (signal.shape[0]-1)/2, num=signal.shape[0])
      y = np.linspace(-(signal.shape[1]-1)/2, (signal.shape[1]-1)/2, num=signal.shape[1])
      f = interpolate.interp2d(x, y, signal)
      x = np.linspace(-(signal.shape[0]-1)/2, (signal.shape[0]-1)/2, num=target_shape[0])
      y = np.linspace(-(signal.shape[1]-1)/2, (signal.shape[1]-1)/2, num=target_shape[1])
      return f(x,y)
  
  
  def calc_energy_and_filters(stim, stim_df, n_orientations=6, save_path_template=None):
      """this creates the energy and filter arrays
  
      We assume the stimuli have natural groups, here indexed by the "class_idx" column in stim_df,
      and all stimuli within these groups should be considered the same stimuli, that is, we sum the
      energy across all of them. for the spatial frequency project, these are the different phases of
      the gratings (because of how we structure our experiment, we estimate a response amplitude to
      all phases together).
  
      Note that this will take a while to run (~10 or 20 minutes). Since it only needs to run once
      per experiment, didn't bother to make it efficient at all. The outputs will also be very large,
      totalling about 11GB
  
      Parameters
      ----------
      stim : np.ndarray
          The stimuli to produce energy for. Should have shape (n, *img_size), where n is the number 
          of total stimuli.
      stim_df : pd.DataFrame
          The DataFrame describing the stimuli. Must contain the column "class_idx", which indexes 
          the different stimulus classes (see above)
      n_orientations : int
          the number of orientations in the steerable pyramid. 6 is the number used to model fMRI
          voxels in Roth, Z. N., Heeger, D., & Merriam, E. (2018). Stimulus vignetting and 
          orientation selectivity in human visual cortex. bioRxiv.
      save_path_template : str or None
          the template string for the save path we'll use for energy and filters. should end in .npy 
          and contain one %s, which we'll replace with "energy" and "filters".
  
      Returns
      -------
      energy : np.ndarray
          energy has shape (stim_df.class_idx.nunique(), max_ht, n_orientations, *img_size) and 
          contains the energy (square and absolute value the complex valued output of 
          SteerablePyramidFreq; equivalently, square and sum the output of the quadrature pair of 
          filters that make up the pyramid) for each image, at each scale and orientation. the energy
          has all been upsampled to the size of the initial image.
      filters : np.ndarray
          filters has shape (max_ht, n_orientations, *img_size) and is the fourier transform of the 
          filters at each scale and orientation, zero-padded so they all have the same size. we only 
          have one set of filters (instead of one per stimulus class) because the same pyramid was 
          used for each of them; we ensure this by getting the filters for each stimulus class and 
          checking that they're individually equal to the average across classes.
  
      """
      img_size = stim.shape[1:]
      # this computation comes from the SteerablePyramidFreq code
      max_ht = int(np.floor(np.log2(min(img_size))) - 2)
      energy = np.zeros((stim_df.class_idx.nunique(), max_ht, n_orientations, *img_size),
                        dtype=np.float32)
      filters = np.zeros_like(energy)
      for i, g in stim_df.groupby('class_idx'):
          idx = g.index
          filled_filters = False
          for j in idx:
              pyr = pt.pyramids.SteerablePyramidFreq(stim[j], order=n_orientations-1, is_complex=True)
              for k, l in itertools.product(range(max_ht), range(n_orientations)):
                  energy[int(i), k, l, :, :] += upsample(np.abs(pyr.pyr_coeffs[(k, l)])**2, img_size)
                  # we only want to run this once per stimulus class
                  if not filled_filters:
                      if k > 0:
                          lomask = pyr._lomasks[k-1]
                      else:
                          lomask = pyr._lo0mask
                      filt = pyr._anglemasks[k][l] * pyr._himasks[k] * lomask
                      pad_num = []
                      for m in range(2):
                          pad_num.append([(img_size[m] - filt.shape[m])//2, (img_size[m] - filt.shape[m])//2])
                          if filt.shape[m] + 2*pad_num[m][0] != img_size[m]:
                              pad_num[m][0] += img_size[m] - (filt.shape[m] + 2*pad_num[m][0])
                      filters[int(i), k, l, :, :] = np.pad(filt, pad_num, 'constant', constant_values=0)
              filled_filters = True
      filter_mean = np.mean(filters, 0)
      for i in range(filters.shape[0]):
          if not(np.allclose(filter_mean, filters[i,:,:,:,:])):
              raise Exception("Something has gone terribly wrong, the filters for stim class %d are different than the rest!" % i)
      filters = filter_mean
      if save_path_template is not None:
          np.save(save_path_template % "energy", energy)
          np.save(save_path_template % "filters", filters)
      return energy, filters
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(
          description=("Calculate and save the energy for each stimulus class, as well as the Fourier"
                       " transform of the filters of the steerable pyramid we use to get this. For "
                       "use with image-computable version of this model"),
          formatter_class=argparse.ArgumentDefaultsHelpFormatter)
      parser.add_argument("stimuli",
                          help=("Path to the stimulus .npy file."))
      parser.add_argument("stimuli_description_df",
                          help=("Path to the stimulus description dataframe .csv file."))
      parser.add_argument("save_path_template",
                          help=("Path template (with .npy extension) where we'll save the results. "
                                "Should contain one %s."))
      parser.add_argument('--n_orientations', '-n', default=6, type=int,
                          help=("The number of orientations in the steerable pyramid used here."))
      args = vars(parser.parse_args())
      stim = np.load(args.pop('stimuli'))
      stim_df = pd.read_csv(args.pop('stimuli_description_df'))
      calc_energy_and_filters(stim, stim_df, **args)
  