text: import pandas as pd
  import itertools as it
  import seaborn as sns
  import numpy as np
  from pymea import matlab_compatibility as mc
  from matplotlib import pyplot as plt
  from matplotlib import mlab as mlab
  import random
  from datetime import datetime, timedelta
  
  def select_neurons(cat_table, min_freq = 0, max_freq = 100000000, data_col='spike_freq'):
      '''
      Returns a version of cat_table only containing units whose mean frequency is between min_freq and max_freq
      '''
      unit_freq_mean = cat_table.groupby(('unit_name'))[data_col].mean()
      unit_freq_mean = unit_freq_mean.rename(data_col).reset_index()
      filt = (unit_freq_mean[data_col] >= min_freq) & (unit_freq_mean[data_col] < max_freq)
      selected_units = unit_freq_mean.loc[filt,'unit_name']
      return (cat_table.loc[cat_table['unit_name'].isin(selected_units)], selected_units)
  
  def filter_neurons_homeostasis(cat_table, baseline_table, stim_table, ind_filter = True, var=10, minHz = 0.001, maxHz = 100000, foldMin = 0.001, filter_wells = False, data_col = 'spike_freq'):
      '''
      Returns a cat_table only including neurons that pass the filters for min/maxHz, baseline var, staying alive
      throughout the experiment, responding to drug, and whose wells are behaving similarly to others.
      '''
      c_filter = pd.DataFrame()
      b_filter = pd.DataFrame()
      count_real = 0
      count_live = 0
      count_final = 0
      last_time = max(cat_table['time'])
      #filter individual neurons based on baseline firing, whether they stay alive, and whether they respond to stim
      for cond in cat_table['condition'].unique():
          c = cat_table.query('condition == "%s"'%cond)
          b = baseline_table.query('condition == "%s"'%cond)
          s = stim_table.query('condition == "%s"'%cond)
          c_filter_cond = pd.DataFrame()
          b_filter_cond = pd.DataFrame()
          s_filter_cond = pd.DataFrame()
          for unit_name in c['unit_name'].unique():
              unit = c.query('unit_name == @unit_name')
              unit_b = b.query('unit_name == @unit_name')
              unit_s = s.query('unit_name == @unit_name')
              meanOfBaseline = np.mean(unit_b[data_col])
              varOfBaseline = (unit_b.loc[:,data_col].max() - unit_b.loc[:,data_col].min())/meanOfBaseline
              meanAfterDrug = np.mean(unit_s[6:42][data_col]) #mean firing in first 3hr of stim, skipping first 30min
              folds = unit[data_col]/meanOfBaseline
              folds_b = unit_b[data_col]/meanOfBaseline
              folds_s = unit_s[data_col]/meanOfBaseline
              if meanOfBaseline > minHz and meanOfBaseline < maxHz and varOfBaseline < var:
                  count_real = count_real+1
                  if (sum((folds) < foldMin)<50):# and (sum(folds > 500)<100):
                      count_live = count_live+1
                      if (ind_filter == False) or (meanOfBaseline < 1*meanAfterDrug): 
                          count_final = count_final+1
                          well_num = mc.get_well_number(unit_name)
                          unit.loc[:,'well'] = well_num
                          unit.loc[:,'folds'] = folds
                          unit_b.loc[:,'well'] = well_num
                          unit_b.loc[:,'folds'] = folds_b
                          unit_s.loc[:,'well'] = well_num
                          unit_s.loc[:,'folds'] = folds_s
                          c_filter_cond = c_filter_cond.append(unit, ignore_index=True)
                          b_filter_cond = b_filter_cond.append(unit_b, ignore_index=True)
                          s_filter_cond = s_filter_cond.append(unit_s, ignore_index=True)
              else:
                  continue
                  
          if c_filter_cond.empty:
              return (c_filter, b_filter, count_real, count_live, count_final)
      
          #filter out wells that have fewer than two neurons, or whose median doesn't increase by at least 1.3
          for w in c_filter_cond['well'].unique():
              s_well = s_filter_cond.query('well == @w')
              num_units = len(s_well['unit_name'].unique())
              if num_units <= 2:
                  c_filter_cond = c_filter_cond[c_filter_cond.well != w]
                  b_filter_cond = b_filter_cond[b_filter_cond.well != w]
                  count_real = count_real - num_units
                  count_live = count_real - num_units
                  count_final = count_final - num_units
                  print('Fewer than 3 neurons. Omitting well ' + repr(w))
              elif ind_filter == True: 
                  well_med = s_well.groupby('time')['folds'].median()
                  fold_med = np.mean(well_med[6:42])
                  if fold_med < 1.3:
                      c_filter_cond = c_filter_cond[c_filter_cond.well != w]
                      b_filter_cond = b_filter_cond[b_filter_cond.well != w]
                      count_final = count_final - num_units
                      print('Weak median FR response. Omitting well ' + repr(w))
          # filter out wells that behave differently from others
          if filter_wells == True:
              if c_filter_cond.empty:
                  return (c_filter_cond, b_filter_cond, count_real, count_live, count_final)
              c_filter_cond, b_filter_cond = drop_outlier_wells(c_filter_cond, b_filter_cond, data_col = data_col)
          
          c_filter = c_filter.append(c_filter_cond, ignore_index=True) #concatenate filtered condition tables
          b_filter = b_filter.append(b_filter_cond, ignore_index=True) #concatenate filtered condition tables
      return (c_filter, b_filter, count_real, count_live, count_final)
  
  def drop_outlier_wells(c_filter, b_filter, data_col = 'spike_freq'):
      '''
      Crops wells whose correlation coefficient with both the mean/median frequency of the entire condition is below 0.6
      '''
      mean_freq_traces = c_filter.groupby(('condition', 'time'))[data_col].mean()
      mean_freq_traces = mean_freq_traces.rename(data_col).reset_index()
      mean_freq_traces_b = b_filter.groupby(('condition', 'time'))[data_col].mean()
      mean_freq_traces_b = mean_freq_traces_b.rename(data_col).reset_index()
      median_freq_traces = c_filter.groupby(('condition', 'time'))[data_col].median()
      median_freq_traces = median_freq_traces.rename(data_col).reset_index()
      median_freq_traces_b = b_filter.groupby(('condition', 'time'))[data_col].median()
      median_freq_traces_b = median_freq_traces_b.rename(data_col).reset_index()
      meanOfMean = np.mean(mean_freq_traces_b[data_col])
      meanOfMedian = np.mean(median_freq_traces_b[data_col])
      mean_folds = mean_freq_traces[data_col]/meanOfMean
      median_folds = median_freq_traces[data_col]/meanOfMedian
      
      mean_by_well = c_filter.groupby(('well', 'time'))[data_col].mean()
      mean_by_well = mean_by_well.rename(data_col).reset_index()
      median_by_well = c_filter.groupby(('well', 'time'))[data_col].median()
      median_by_well = median_by_well.rename(data_col).reset_index()
      mean_by_well_b = b_filter.groupby(('well', 'time'))[data_col].mean()
      mean_by_well_b = mean_by_well_b.rename(data_col).reset_index()
      median_by_well_b = b_filter.groupby(('well', 'time'))[data_col].median()
      median_by_well_b = median_by_well_b.rename(data_col).reset_index()
  
      for well in mean_by_well['well'].unique():
          this_well_mean = mean_by_well.query('well == @well')
          this_well_mean_b = mean_by_well_b.query('well == @well')
          this_well_median = median_by_well.query('well == @well')
          this_well_median_b = median_by_well_b.query('well == @well')
          meanOfMean_well = np.mean(this_well_mean_b[data_col])
          meanOfMedian_well = np.mean(this_well_median_b[data_col])
          
          mean_folds_well = this_well_mean[data_col]/meanOfMean_well
          median_folds_well = this_well_median[data_col]/meanOfMedian_well
          corr_mean = np.corrcoef(mean_folds_well, mean_folds)
          corr_median = np.corrcoef(median_folds_well, median_folds)
          #print(corr_mean)
          #print(corr_median)
          if corr_mean[0,1] < 0.6 and corr_median[0,1] < 0.6:
              c_filter = c_filter[c_filter.well != well]
              b_filter = b_filter[b_filter.well != well]
              print('Omitting well ' + repr(well))
      return (c_filter, b_filter)
  
  def cdf(data):
      '''
      returns a sorted version of data (small to big) and an array of their proportions for a cdf plot
      '''
      sorted_data = np.sort(data)
      # make array of proportions from 0:1, the length of the data
      p = 1. * np.arange(len(data)) / (len(data) - 1)
      return (sorted_data, p)
  
  def select_homeo_units(plot_group, c_filter, b_filter):
      ''' 
      Returns the data table of neurons whose end firing is below (plot_group=1) or above (plot_group=2) 
      double the baseline firing. Used to show plots with only the neurons that homeostased or didn't.
      '''
      low_units = []
      mid_units = []
      high_units = []
      baseline_stop = max(b_filter['time'])
      end_time = max(c_filter['time']) - timedelta(hours = 1)
      lasthr_drug = end_time - timedelta(hours = 1)
      lasthr_baseline = baseline_stop - timedelta(hours = 1)
          
      for unit in c_filter['unit_name'].unique():
          unit_end = c_filter.query('unit_name == @unit and time >= @lasthr_drug and time <= @end_time')
          unit_start = c_filter.query('unit_name == @unit and time >= @lasthr_baseline and time < @baseline_stop')
          mean_end = unit_end['spike_freq'].mean()
          mean_start = unit_start['spike_freq'].mean()
          homeo_val = mean_end/mean_start
          if homeo_val < 0.5:
              low_units.append(unit)
          elif homeo_val <= 1.5:
              mid_units.append(unit)
          else:
              high_units.append(unit)
      if plot_group == 1:
          # select neurons that died
          c_filter = c_filter.loc[c_filter['unit_name'].isin(low_units)]
          b_filter = b_filter.loc[b_filter['unit_name'].isin(low_units)]
      elif plot_group == 2:
          # select neurons that worked
          c_filter = c_filter.loc[c_filter['unit_name'].isin(mid_units)]
          b_filter = b_filter.loc[b_filter['unit_name'].isin(mid_units)]
      elif plot_group == 3:
          # select neurons that never decreased their firing
          c_filter = c_filter.loc[c_filter['unit_name'].isin(high_units)]
          b_filter = b_filter.loc[b_filter['unit_name'].isin(high_units)]
      return (c_filter, b_filter)
      
  def calc_end_vs_start(units, baseline_stop, cat_table, end_time = 0):
      '''
      Calculates the ratio of FR at the end of the experiment to pre-drug FR.
      '''
      if end_time == 0:
          end_time = max(cat_table['time'])
      lasthr_drug = end_time - timedelta(hours = 1)
      lasthr_baseline = baseline_stop - timedelta(hours = 1)
      
      units_table = cat_table.loc[cat_table['unit_name'].isin(units)]
      units_end = units_table.query('time >= @lasthr_drug')
      units_start = units_table.query('time >= @lasthr_baseline and time < @baseline_stop')
      mean_end = units_end.groupby(('unit_name'))['spike_freq'].mean()
      mean_end = mean_end.rename('spike_freq').reset_index()
      mean_start = units_start.groupby(('unit_name'))['spike_freq'].mean()
      mean_start = mean_start.rename('spike_freq').reset_index()
      ratio = mean_end['spike_freq']/mean_start['spike_freq']
      return pd.DataFrame(data = {'unit_name': mean_end['unit_name'], 'ratio': ratio})