notebooks to docs, small fixes for plots

Author: mmyros

bayes_window/utils.py

@@ -16,13 +16,14 @@ def add_data_to_posterior(df,
                           conditions=None,  # eg ('stim_on', 'stim_stop')
                           b_name='b_stim_per_condition',  # for posterior
                           group_name='Condition code',  # for posterior
-                          do_make_change=True,
+                          do_make_change='subtract',
                           do_mean_over_trials=True,
                           ):
     index_cols = list(index_cols)
-    if conditions is None:
-        conditions = df[condition_name].drop_duplicates().sort_values().values
+
+    conditions= conditions or df[condition_name].drop_duplicates().sort_values().values
     assert len(conditions) == 2, f'{condition_name}={conditions}. Should be only two instead!'
+    assert do_make_change in [False, 'subtract', 'divide']
     if not (condition_name in index_cols):
         index_cols.append(condition_name)
     if do_mean_over_trials:
@@ -36,6 +37,7 @@ def add_data_to_posterior(df,
                                  index_cols=index_cols,
                                  condition_name=condition_name,
                                  conditions=conditions,
+                                 fold_change_method=do_make_change,
                                  do_take_mean=False)
         # Condition is removed from both index columns and dfbayes
         index_cols.remove(condition_name)
@@ -98,29 +100,33 @@ def fill_row(rows):
 
 
 def make_fold_change(df, y='log_firing_rate', index_cols=('Brain region', 'Stim phase'),
-                     condition_name='stim', conditions=(0, 1), do_take_mean=False):
-    # for index_col in index_cols:
-    #     assert type(df[index_col].iloc[0]) != str, f'Make sure {index_col} contains not strings!'
+                     condition_name='stim', conditions=(0, 1), do_take_mean=False, fold_change_method='divide'):
     for condition in conditions:
         assert condition in df[condition_name].unique(), f'{condition} not in {df[condition_name].unique()}'
     if y not in df.columns:
         raise ValueError(f'{y} is not a column in this dataset: {df.columns}')
+
+    # Take mean of trials:
     if do_take_mean:
-        # Take mean of trials:
         df = df.groupby(list(index_cols)).mean().reset_index()
+
     # Make multiindex
     mdf = df.set_index(list(set(index_cols) - {'i_spike'})).copy()
-    # mdf.xs(0, level='stim') - mdf.xs(1, level='stim')
     if (mdf.xs(conditions[1], level=condition_name).size !=
         mdf.xs(conditions[0], level=condition_name).size):
         raise IndexError(f'Uneven number of entries in conditions! Try setting do_take_mean=True'
                          f'{mdf.xs(conditions[0], level=condition_name).size, mdf.xs(conditions[1], level=condition_name).size}')
 
     # Subtract/divide
     try:
-        data = (mdf.xs(conditions[1], level=condition_name) -
-                mdf.xs(conditions[0], level=condition_name)
-                ).reset_index()
+        if fold_change_method == 'subtract':
+            data = (mdf.xs(conditions[1], level=condition_name) -
+                    mdf.xs(conditions[0], level=condition_name)
+                    ).reset_index()
+        else:
+            data = (mdf.xs(conditions[1], level=condition_name) /
+                    mdf.xs(conditions[0], level=condition_name)
+                    ).reset_index()
     except Exception as e:
         print(f'Try recasting {condition_name} as integer and try again. Alternatively, use bayes_window.workflow.'
               f' We do that automatically there ')

bayes_window/visualization.py

@@ -13,7 +13,6 @@ def facet(base_chart,
           width=80,
           height=150,
           ):
-    print('custom facet')
     alt.themes.enable('vox')
     if column is None and row is None:
         raise RuntimeError('Need either column, or row, or both!')
@@ -67,32 +66,18 @@ def plot_data(df=None, x=None, y=None, color=None, add_box=True, base_chart=None
     chart = base.mark_line(fill=None, opacity=.5, size=3).encode(
         x=x,
         color=f'{color}',
-        y=f'{y}:Q'
+        y=alt.Y(f'{y}:Q')
     )
     if add_box:
         # Shift x axis for box so that it doesnt overlap:
         # df['x_box'] = df[x[:-2]] + .01
         chart += base.mark_boxplot(opacity=.3, size=12, color='black').encode(
             x=x,
-            y=f'{y}:Q'
+            y=alt.Y(f'{y}:Q')
         )
     return chart
 
 
-# def plot_data_and_posterior(df, y='Coherence diff', title='coherence', x='Stim phase', color='Subject',
-#                             add_box=True, **kwargs):
-#     # Keep kwargs!
-#     assert (x in df) | (x[:-2] in df), f'Column {x} is not present in data: {df.columns}'
-#     assert color in df
-#     assert y in df.columns, f'{y} is not in {df.columns}'
-#
-#     chart_d = plot_data(df=df, x=x, y=y, color=color, add_box=add_box, base_chart=alt.Chart(df))
-#     chart_p = plot_posterior(df, title=title, x=x, base_chart=alt.Chart(df))
-#     chart = chart_d + chart_p
-#
-#     return chart
-
-
 def plot_posterior(df=None, title='', x='Stim phase', do_make_change=True, base_chart=None, **kwargs):
     assert (df is not None) or (base_chart is not None)
     data = base_chart.data if df is None else df
@@ -244,15 +229,15 @@ def fake_spikes_explore(df, df_monster, index_cols):
         y=alt.Y(y, scale=alt.Scale(zero=True)),
     ).properties(width=width, height=240).facet(
         # row='mouse:N',
-        column=alt.Column('mouse'))  # .resolve_scale(y='independent')
+        column=alt.Column('mouse'))
 
     bar = (alt.Chart(data=data_fold_change).mark_bar().encode(y=alt.Y(y, aggregate='mean')) +
            alt.Chart(data=data_fold_change).mark_errorbar().encode(y=alt.Y(y, aggregate='stderr'))).encode(
         x=alt.X('neuron:N', ),
         y=alt.Y(y),
     ).properties(width=width * 2, height=240).facet(
         # row='Inversion:N',
-        column=alt.Column('mouse'))  # .resolve_scale(y='independent')
+        column=alt.Column('mouse'))
 
     bar_combined = (alt.Chart(data=data_fold_change).mark_bar().encode(y=alt.Y(y, aggregate='mean')) +
                     alt.Chart(data=data_fold_change).mark_errorbar().encode(y=alt.Y(y, aggregate='stderr'))).encode(

bayes_window/workflow.py

@@ -2,12 +2,13 @@
 from importlib import reload
 
 import altair as alt
+from sklearn.preprocessing import LabelEncoder
+
 from bayes_window import models
 from bayes_window import utils
 from bayes_window import visualization
 from bayes_window.fitting import fit_numpyro
 from bayes_window.visualization import plot_posterior
-from sklearn.preprocessing import LabelEncoder
 
 reload(visualization)
 reload(utils)
@@ -54,9 +55,11 @@ def fit_conditions(self, model=models.model_single_lognormal, add_data=True):
                                                                   do_make_change=False
                                                                   )
 
-    def fit_slopes(self, add_data=True, model=models.model_hier_normal_stim,
+    def fit_slopes(self, add_data=True, model=models.model_hier_normal_stim, do_make_change='subtract',
                    plot_index_cols=None):
+        assert do_make_change in ['subtract', 'divide']
         self.bname = 'b_stim_per_condition'
+        self.do_make_change = do_make_change
         if plot_index_cols is None:
             plot_index_cols = self.levels[:3]
         # By convention, top condition is first in list of levels:
@@ -89,7 +92,7 @@ def fit_slopes(self, add_data=True, model=models.model_hier_normal_stim,
                                                     condition_name=top_condition,
                                                     b_name=self.bname,
                                                     group_name=self.levels[-1],
-                                                    do_make_change=True,
+                                                    do_make_change=do_make_change,
                                                     do_mean_over_trials=True,
                                                     )
         else:  # Just convert posterior to dataframe
@@ -107,18 +110,21 @@ def plot_posteriors_slopes(self, x=None, color=None, add_box=True, independent_a
         # Plot posterior
         if hasattr(self, 'data_and_posterior'):
             base_chart = alt.Chart(self.data_and_posterior)
-            chart_p = plot_posterior(title=f'd_{self.y}', x=x, base_chart=base_chart)
+            chart_p = plot_posterior(title=f'{self.y}', x=x, base_chart=base_chart)
         else:
             base_chart = alt.Chart(self.data)
             add_data = True  # Otherwise nothing to do
 
         if add_data:
             chart_d = visualization.plot_data(x=x, y=f'{self.y} diff', color=color, add_box=add_box,
                                               base_chart=base_chart)
-            self.chart = chart_p + chart_d
+            self.chart = chart_d + chart_p  # Not chart_d + chart_p, or bayes means and HPD get scaled independently
         else:
             self.chart = chart_p
-        assert self.chart.data is not None
+        if independent_axes:
+            self.chart = self.chart.resolve_scale(y='independent')
+        elif self.do_make_change == 'divide':
+            warnings.warn('division change and independent axes will lead to separate axis! Upstream bug I think')
         return self.chart
 
     # TODO plot_posteriors_slopes and plot_posteriors_no_slope can be one