Merge branch 'develop' into single_source_version

Author: nbonacchi

MANIFEST.in

@@ -1,5 +1,6 @@
 include ibllib/atlas/allen_structure_tree.csv
 include ibllib/atlas/beryl.npy
+include ibllib/atlas/cosmos.npy
 include ibllib/io/extractors/extractor_types.json
 include brainbox/tests/wheel_test.p
 recursive-include brainbox/tests/fixtures *

brainbox/behavior/training.py

@@ -1,12 +1,11 @@
-import logging
 from one.api import ONE
 import datetime
 import re
 import numpy as np
 from iblutil.util import Bunch
 import brainbox.behavior.pyschofit as psy
-
-logger = logging.getLogger('ibllib')
+import logging
+_logger = logging.getLogger('ibllib')
 
 
 def get_lab_training_status(lab, date=None, details=True, one=None):
@@ -110,7 +109,7 @@ def get_sessions(subj, date=None, one=None):
 
         # If still 0 sessions then return with warning
         if len(sessions) == 0:
-            logger.warning(f"No training sessions detected for {subj}")
+            _logger.warning(f"No training sessions detected for {subj}")
             return [None] * 4
 
     trials = Bunch()
@@ -274,30 +273,30 @@ def display_status(subj, sess_dates, status, perf_easy=None, n_trials=None, psyc
     """
 
     if perf_easy is None:
-        logger.info(f"\n{subj} : {status} \nSession dates=[{sess_dates[0]}, {sess_dates[1]}, "
-                    f"{sess_dates[2]}]")
+        print(f"\n{subj} : {status} \nSession dates=[{sess_dates[0]}, {sess_dates[1]}, "
+              f"{sess_dates[2]}]")
     elif psych_20 is None:
-        logger.info(f"\n{subj} : {status} \nSession dates={[x for x in sess_dates]}, "
-                    f"Perf easy={[np.around(pe,2) for pe in perf_easy]}, "
-                    f"N trials={[nt for nt in n_trials]} "
-                    f"\nPsych fit over last 3 sessions: "
-                    f"bias={np.around(psych[0],2)}, thres={np.around(psych[1],2)}, "
-                    f"lapse_low={np.around(psych[2],2)}, lapse_high={np.around(psych[3],2)} "
-                    f"\nMedian reaction time at 0 contrast over last 3 sessions = "
-                    f"{np.around(rt,2)}")
+        print(f"\n{subj} : {status} \nSession dates={[x for x in sess_dates]}, "
+              f"Perf easy={[np.around(pe,2) for pe in perf_easy]}, "
+              f"N trials={[nt for nt in n_trials]} "
+              f"\nPsych fit over last 3 sessions: "
+              f"bias={np.around(psych[0],2)}, thres={np.around(psych[1],2)}, "
+              f"lapse_low={np.around(psych[2],2)}, lapse_high={np.around(psych[3],2)} "
+              f"\nMedian reaction time at 0 contrast over last 3 sessions = "
+              f"{np.around(rt,2)}")
 
     else:
-        logger.info(f"\n{subj} : {status} \nSession dates={[x for x in sess_dates]}, "
-                    f"Perf easy={[np.around(pe,2) for pe in perf_easy]}, "
-                    f"N trials={[nt for nt in n_trials]} "
-                    f"\nPsych fit over last 3 sessions (20): "
-                    f"bias={np.around(psych_20[0],2)}, thres={np.around(psych_20[1],2)}, "
-                    f"lapse_low={np.around(psych_20[2],2)}, lapse_high={np.around(psych_20[3],2)} "
-                    f"\nPsych fit over last 3 sessions (80): bias={np.around(psych_80[0],2)}, "
-                    f"thres={np.around(psych_80[1],2)}, lapse_low={np.around(psych_80[2],2)}, "
-                    f"lapse_high={np.around(psych_80[3],2)} "
-                    f"\nMedian reaction time at 0 contrast over last 3 sessions = "
-                    f"{np.around(rt, 2)}")
+        print(f"\n{subj} : {status} \nSession dates={[x for x in sess_dates]}, "
+              f"Perf easy={[np.around(pe,2) for pe in perf_easy]}, "
+              f"N trials={[nt for nt in n_trials]} "
+              f"\nPsych fit over last 3 sessions (20): "
+              f"bias={np.around(psych_20[0],2)}, thres={np.around(psych_20[1],2)}, "
+              f"lapse_low={np.around(psych_20[2],2)}, lapse_high={np.around(psych_20[3],2)} "
+              f"\nPsych fit over last 3 sessions (80): bias={np.around(psych_80[0],2)}, "
+              f"thres={np.around(psych_80[1],2)}, lapse_low={np.around(psych_80[2],2)}, "
+              f"lapse_high={np.around(psych_80[3],2)} "
+              f"\nMedian reaction time at 0 contrast over last 3 sessions = "
+              f"{np.around(rt, 2)}")
 
 
 def concatenate_trials(trials):

brainbox/ephys_plots.py

@@ -1,9 +1,10 @@
 import numpy as np
 from matplotlib import cm
-
+import matplotlib.pyplot as plt
 from brainbox.plot_base import (ImagePlot, ScatterPlot, ProbePlot, LinePlot, plot_line,
                                 plot_image, plot_probe, plot_scatter, arrange_channels2banks)
 from brainbox.processing import bincount2D, compute_cluster_average
+from ibllib.atlas.regions import BrainRegions
 
 
 def image_lfp_spectrum_plot(lfp_power, lfp_freq, chn_coords, chn_inds, freq_range=(0, 300),
@@ -372,3 +373,102 @@ def line_amp_plot(spike_amps, spike_depths, spike_times, chn_coords, d_bin=10, d
         fig, ax = plot_line(data.convert2dict())
         return data.convert2dict(), fig, ax
     return data
+
+
+def plot_brain_regions(channel_ids, channel_depths=None, brain_regions=None, display=True, ax=None):
+    """
+    Plot brain regions along probe, if channel depths is provided will plot along depth otherwise along channel idx
+    :param channel_ids: atlas ids for each channel
+    :param channel_depths: depth along probe for each channel
+    :param brain_regions: BrainRegions object
+    :param display: whether to output plot
+    :param ax: axis to plot on
+    :return:
+    """
+
+    if channel_depths is not None:
+        assert channel_ids.shape[0] == channel_depths.shape[0]
+
+    br = brain_regions or BrainRegions()
+
+    region_info = br.get(channel_ids)
+    boundaries = np.where(np.diff(region_info.id) != 0)[0]
+    boundaries = np.r_[0, boundaries, region_info.id.shape[0] - 1]
+
+    regions = np.c_[boundaries[0:-1], boundaries[1:]]
+    if channel_depths is not None:
+        regions = channel_depths[regions]
+    region_labels = np.c_[np.mean(regions, axis=1), region_info.acronym[boundaries[1:]]]
+    region_colours = region_info.rgb[boundaries[1:]]
+
+    if display:
+        if ax is None:
+            fig, ax = plt.subplots()
+        else:
+            fig = ax.get_figure()
+
+        for reg, col in zip(regions, region_colours):
+            height = np.abs(reg[1] - reg[0])
+            color = col / 255
+            ax.bar(x=0.5, height=height, width=1, color=color, bottom=reg[0], edgecolor='w')
+        ax.set_yticks(region_labels[:, 0].astype(int))
+        ax.yaxis.set_tick_params(labelsize=8)
+        ax.get_xaxis().set_visible(False)
+        ax.set_yticklabels(region_labels[:, 1])
+        ax.spines['right'].set_visible(False)
+        ax.spines['top'].set_visible(False)
+        ax.spines['bottom'].set_visible(False)
+
+        return fig, ax
+    else:
+        return regions, region_labels, region_colours
+
+
+def plot_cdf(spike_amps, spike_depths, spike_times, n_amp_bins=10, d_bin=40, amp_range=None, d_range=None,
+             display=False, cmap='hot'):
+    """
+    Plot cumulative amplitude of spikes across depth
+    :param spike_amps:
+    :param spike_depths:
+    :param spike_times:
+    :param n_amp_bins: number of amplitude bins to use
+    :param d_bin: the value of the depth bins in um (default is 40 um)
+    :param amp_range: amp range to use [amp_min, amp_max], if not given automatically computed from spike_amps
+    :param d_range: depth range to use, by default [0, 3840]
+    :param display: whether or not to display plot
+    :param cmap:
+    :return:
+    """
+
+    amp_range = amp_range or np.quantile(spike_amps, (0, 0.9))
+    amp_bins = np.linspace(amp_range[0], amp_range[1], n_amp_bins)
+    d_range = d_range or [0, 3840]
+    depth_bins = np.arange(d_range[0], d_range[1] + d_bin, d_bin)
+    t_bin = np.max(spike_times)
+
+    def histc(x, bins):
+        map_to_bins = np.digitize(x, bins)  # Get indices of the bins to which each value in input array belongs.
+        res = np.zeros(bins.shape)
+
+        for el in map_to_bins:
+            res[el - 1] += 1  # Increment appropriate bin.
+        return res
+
+    cdfs = np.empty((len(depth_bins) - 1, n_amp_bins))
+    for d in range(len(depth_bins) - 1):
+        spikes = np.bitwise_and(spike_depths > depth_bins[d], spike_depths <= depth_bins[d + 1])
+        h = histc(spike_amps[spikes], amp_bins) / t_bin
+        hcsum = np.cumsum(h[::-1])
+        cdfs[d, :] = hcsum[::-1]
+
+    cdfs[cdfs == 0] = np.nan
+
+    data = ImagePlot(cdfs.T, x=amp_bins * 1e6, y=depth_bins[:-1], cmap=cmap)
+    data.set_labels(title='Cumulative Amplitude', xlabel='Spike amplitude (uV)',
+                    ylabel='Distance from probe tip (um)', clabel='Firing Rate (Hz)')
+
+    if display:
+        fig, ax = plot_image(data.convert2dict(), fig_kwargs={'figsize': [3, 7]})
+        return data.convert2dict(), fig, ax
+
+    return data