Merge pull request #103 from JaerongA/no_curation_plot

add ephys_report schema for data visualizations

Author: Thinh Nguyen

.gitignore

@@ -118,4 +118,7 @@ dj_local_conf_old.json
 docker-compose.y*ml
 
 # include
-!docs/docker-compose.yaml
+!docs/docker-compose.yaml
+
+# vscode settings
+.vscode

element_array_ephys/__init__.py

@@ -3,12 +3,13 @@
 import os
 
 
-dj.config['enable_python_native_blobs'] = True
+dj.config["enable_python_native_blobs"] = True
 
 
 def get_logger(name):
     log = logging.getLogger(name)
-    log.setLevel(os.getenv('LOGLEVEL', 'INFO'))
+    log.setLevel(os.getenv("LOGLEVEL", "INFO"))
     return log
 
+
 from . import ephys_acute as ephys

element_array_ephys/ephys_acute.py

@@ -0,0 +1,157 @@
+import pathlib
+import datetime
+import datajoint as dj
+import typing as T
+import json
+
+schema = dj.schema()
+
+ephys = None
+
+
+def activate(schema_name, ephys_schema_name, *, create_schema=True, create_tables=True):
+    """
+    activate(schema_name, *, create_schema=True, create_tables=True, activated_ephys=None)
+        :param schema_name: schema name on the database server to activate the `ephys_report` schema
+        :param ephys_schema_name: schema name of the activated ephys element for which this ephys_report schema will be downstream from
+        :param create_schema: when True (default), create schema in the database if it does not yet exist.
+        :param create_tables: when True (default), create tables in the database if they do not yet exist.
+    (The "activation" of this ephys_report module should be evoked by one of the ephys modules only)
+    """
+    global ephys
+    ephys = dj.create_virtual_module("ephys", ephys_schema_name)
+    schema.activate(
+        schema_name,
+        create_schema=create_schema,
+        create_tables=create_tables,
+        add_objects=ephys.__dict__,
+    )
+
+
+@schema
+class ProbeLevelReport(dj.Computed):
+    definition = """
+    -> ephys.CuratedClustering
+    shank         : tinyint unsigned
+    ---
+    drift_map_plot: attach
+    """
+
+    def make(self, key):
+
+        from . import probe
+        from .plotting.probe_level import plot_driftmap
+
+        save_dir = _make_save_dir()
+
+        units = ephys.CuratedClustering.Unit & key & "cluster_quality_label='good'"
+
+        shanks = set((probe.ProbeType.Electrode & units).fetch("shank"))
+
+        for shank_no in shanks:
+
+            table = (
+                units
+                * ephys.ProbeInsertion.proj()
+                * probe.ProbeType.Electrode.proj("shank")
+                & {"shank": shank_no}
+            )
+
+            spike_times, spike_depths = table.fetch(
+                "spike_times", "spike_depths", order_by="unit"
+            )
+
+            # Get the figure
+            fig = plot_driftmap(spike_times, spike_depths, colormap="gist_heat_r")
+            fig_prefix = (
+                "-".join(
+                    [
+                        v.strftime("%Y%m%d%H%M%S")
+                        if isinstance(v, datetime.datetime)
+                        else str(v)
+                        for v in key.values()
+                    ]
+                )
+                + f"-{shank_no}"
+            )
+
+            # Save fig and insert
+            fig_dict = _save_figs(
+                figs=(fig,),
+                fig_names=("drift_map_plot",),
+                save_dir=save_dir,
+                fig_prefix=fig_prefix,
+                extension=".png",
+            )
+
+            self.insert1({**key, **fig_dict, "shank": shank_no})
+
+
+@schema
+class UnitLevelReport(dj.Computed):
+    definition = """
+    -> ephys.CuratedClustering.Unit
+    ---
+    cluster_quality_label   : varchar(100) 
+    waveform_plotly         : longblob  
+    autocorrelogram_plotly  : longblob
+    depth_waveform_plotly   : longblob
+    """
+
+    def make(self, key):
+
+        from .plotting.unit_level import (
+            plot_waveform,
+            plot_correlogram,
+            plot_depth_waveforms,
+        )
+
+        sampling_rate = (ephys.EphysRecording & key).fetch1(
+            "sampling_rate"
+        ) / 1e3  # in kHz
+
+        peak_electrode_waveform, spike_times, cluster_quality_label = (
+            (ephys.CuratedClustering.Unit & key) * ephys.WaveformSet.PeakWaveform
+        ).fetch1("peak_electrode_waveform", "spike_times", "cluster_quality_label")
+
+        # Get the figure
+        waveform_fig = plot_waveform(
+            waveform=peak_electrode_waveform, sampling_rate=sampling_rate
+        )
+
+        correlogram_fig = plot_correlogram(
+            spike_times=spike_times, bin_size=0.001, window_size=1
+        )
+
+        depth_waveform_fig = plot_depth_waveforms(ephys, unit_key=key, y_range=60)
+
+        self.insert1(
+            {
+                **key,
+                "cluster_quality_label": cluster_quality_label,
+                "waveform_plotly": waveform_fig.to_plotly_json(),
+                "autocorrelogram_plotly": correlogram_fig.to_plotly_json(),
+                "depth_waveform_plotly": depth_waveform_fig.to_plotly_json(),
+            }
+        )
+
+
+def _make_save_dir(root_dir: pathlib.Path = None) -> pathlib.Path:
+    if root_dir is None:
+        root_dir = pathlib.Path().absolute()
+    save_dir = root_dir / "ephys_figures"
+    save_dir.mkdir(parents=True, exist_ok=True)
+    return save_dir
+
+
+def _save_figs(
+    figs, fig_names, save_dir, fig_prefix, extension=".png"
+) -> T.Dict[str, pathlib.Path]:
+    fig_dict = {}
+    for fig, fig_name in zip(figs, fig_names):
+        fig_filepath = save_dir / (fig_prefix + "_" + fig_name + extension)
+        fig.tight_layout()
+        fig.savefig(fig_filepath)
+        fig_dict[fig_name] = fig_filepath.as_posix()
+
+    return fig_dict

element_array_ephys/ephys_chronic.py

@@ -0,0 +1,84 @@
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+
+def plot_raster(units, spike_times) -> matplotlib.figure.Figure:
+
+    units = np.arange(1, len(units) + 1)
+    x = np.hstack(spike_times)
+    y = np.hstack([np.full_like(s, u) for u, s in zip(units, spike_times)])
+    fig, ax = plt.subplots(1, 1, figsize=(32, 8), dpi=100)
+    ax.plot(x, y, "|")
+    ax.set(
+        xlabel="Time (s)",
+        ylabel="Unit",
+        xlim=[0 - 0.5, x[-1] + 0.5],
+        ylim=(1, len(units)),
+    )
+    sns.despine()
+    fig.tight_layout()
+
+    return fig
+
+
+def plot_driftmap(
+    spike_times: np.ndarray, spike_depths: np.ndarray, colormap="gist_heat_r"
+) -> matplotlib.figure.Figure:
+
+    spike_times = np.hstack(spike_times)
+    spike_depths = np.hstack(spike_depths)
+
+    # Time-depth 2D histogram
+    time_bin_count = 1000
+    depth_bin_count = 200
+
+    spike_bins = np.linspace(0, spike_times.max(), time_bin_count)
+    depth_bins = np.linspace(0, np.nanmax(spike_depths), depth_bin_count)
+
+    spk_count, spk_edges, depth_edges = np.histogram2d(
+        spike_times, spike_depths, bins=[spike_bins, depth_bins]
+    )
+    spk_rates = spk_count / np.mean(np.diff(spike_bins))
+    spk_edges = spk_edges[:-1]
+    depth_edges = depth_edges[:-1]
+
+    # Canvas setup
+    fig = plt.figure(figsize=(12, 5), dpi=200)
+    grid = plt.GridSpec(15, 12)
+
+    ax_cbar = plt.subplot(grid[0, 0:10])
+    ax_driftmap = plt.subplot(grid[2:, 0:10])
+    ax_spkcount = plt.subplot(grid[2:, 10:])
+
+    # Plot main
+    im = ax_driftmap.imshow(
+        spk_rates.T,
+        aspect="auto",
+        cmap=colormap,
+        extent=[spike_bins[0], spike_bins[-1], depth_bins[-1], depth_bins[0]],
+    )
+    # Cosmetic
+    ax_driftmap.invert_yaxis()
+    ax_driftmap.set(
+        xlabel="Time (s)",
+        ylabel="Distance from the probe tip ($\mu$m)",
+        ylim=[depth_edges[0], depth_edges[-1]],
+    )
+
+    # Colorbar for firing rates
+    cb = fig.colorbar(im, cax=ax_cbar, orientation="horizontal")
+    cb.outline.set_visible(False)
+    cb.ax.xaxis.tick_top()
+    cb.set_label("Firing rate (Hz)")
+    cb.ax.xaxis.set_label_position("top")
+
+    # Plot spike count
+    ax_spkcount.plot(spk_count.sum(axis=0) / 10e3, depth_edges, "k")
+    ax_spkcount.set_xlabel("Spike count (x$10^3$)")
+    ax_spkcount.set_yticks([])
+    ax_spkcount.set_ylim(depth_edges[0], depth_edges[-1])
+    sns.despine()
+
+    return fig