# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Snakemake Introduction to Elephant
# ## Workflow management based on a electrophysiology example

# <table  bgcolor="#000000"><tr>
# <td><img src=logos/snakemake_logo.svg alt="Drawing" style="width: 300px;"/>
# <td><img src=http://neo.readthedocs.org/en/latest/_images/neologo.png alt="Drawing" style="width: 400px;"/>
# <td><img src=http://elephant.readthedocs.org/en/latest/_static/elephant_logo_sidebar.png alt="Drawing" style="width: 400px;"/>
# </tr></table>

# In this tutorial I will first introduce the workflow management system [snakemake](https://snakemake.readthedocs.io) and show how this can be used for analysis of electrophysiology data using the electrophysiology analysis toolkit [Elephant](http://neuralensemble.org/elephant). For snakemake as well as Elephant tutorials exist already as part of the documentation and were partially reused in the tutorial presented here.

# *NOTE: When running this tutorial locally, please first install the [requirements](environment.yml) to ensure the examples are working.*

# ## Workflow management - Do I need this?

# There are a number of reasons for managing your workflows besides the classical 'I always run first script a and then script b':
# * **Growing Complexity** When starting a seemingly small and simple project everything is still pretty easy, but typically projects tend to grow beyond the initially expected size and complexity. This affects two different aspects 1) the dependency between different steps in workflow as well as 2) the dependency of your analysis results on versions of intermediate steps in the pipeline (data as well as code).
# * **Collaboration & Sharing** Some day a colleague of yours wants to use your workflow or you are required to publish the analysis together with the manuscript presenting the results of your work to the scientific community. Instead of writing a book about which versions of what program you installed in which order, wouldn't it be nice to have a structured and self explanatory compilation of your analysis steps already at hand?
# * **Software Evolution** Luckily, software develops and bugs are fixed from time to time. Unfortunately this also implies that your workflow might break unexpectedly upon updating your system. Having a well defined environment to run your analysis in is essential for reproducibility of your results.

# ## Snakemake  <img src=logos/snakemake_logo.svg alt="Drawing" style="width: 100px;"/>

# Snakemake helps you to structure your workflow by providing a framework to specify the dependencies between individual steps in your analysis workflow. By doing so it enforces a modular structure in the project and allows to specify the (Python) software versions used in each step of the process.

# The workflow definition according to snakemake is file based, i.e. each step in the workflow is specified via the required input files and the generated output files, as you might know from common [`Makefiles`](https://en.wikipedia.org/wiki/Make_(software)). Each step is defined in a `rule`, specifying input and output files as well as instructions on how to get from the input to the output files. Here the instructions to convert `fileA.txt` into `fileB.txt` are a simple copy performed in a shell:

# %%writefile Snakefile
rule:
    input: 'fileA.txt'
    output: 'fileB.txt'
    shell: 'cp fileA.txt fileB.txt'

# For snakemake the workflow definition needs to be specified in a `Snakefile` and can be executed by calling `snakemake` in a terminal in the same location as the `Snakefile`. Here the example rule above has been exported into a [`Snakefile`](Snakefile) using the `%%writefile` jupyter magic command.

# Now we can ask snakemake to generate `fileB.txt` for us:

# + language="sh"
# snakemake fileB.txt
# -

# This fails with snakemake complaining about
# ```
# Missing input files for rule 1:
# fileA.txt
# ```
# Which is correct, since there is no `fileA.txt` present to generate `fileB.txt` from. So let's add second rule which is capable of generating `fileA.txt` without required inputfiles.
#

# %%writefile -a Snakefile
rule:
    output: 'fileA.txt'
    shell: 'touch fileA.txt'

# And ask snakemake again to generate the `fileB.txt` for us

# + language="sh"
# snakemake
# -

# Internally snakemake is first resolving the set of rules into a directed acyclic graph (dag) to determine in which order the rules need to be executed. We can generate a visualization of the workflow using the `--dag` flag in combination with `dot` and `display` (for local notebook instances) or save the graph as svg (e.g. for remote instances).

# + language="sh"
# snakemake --dag | dot | display
# snakemake --dag | dot -Tsvg > dag0.svg
# -

# The resulting graph shows the dependencies between the two rules, which were automatically enumerated. The line style (continuous/dashed) indicated whether the rules were already executed or not.

# ![DAG](dag0.svg)

# We can also provide explicit names for rules to make the graph more human readable:

# %%writefile Snakefile
rule copy_A_to_B:
    input: 'fileA.txt'
    output: 'fileB.txt'
    shell: 'cp {input} {output}'
rule create_A:
    output: 'fileA.txt'
    shell: 'touch fileA.txt'

# + language="sh"
# snakemake --dag | dot | display
# snakemake --dag | dot -Tsvg > dag1.svg
# -

# ![DAG](dag1.svg)

# Here we already used a different notation to specify in the shell command `cp {input} {output}` instead of explicitly repeating the input and output filenames. These placeholders will be substituted by snakemake during execution by the filenames defined as `input` / `output`. We can use the same notation to generalize the required input of the rule depending on the output, e.g., we permit the copy rule to work for arbitrary files having a certain naming scheme. Here a new folder `new_folder` is automatically generated for the copied files.

# +
# %%writefile Snakefile

rule copy_to_new_folder:
    input: 'original_data/file{id}.txt'
    output: 'new_data/file{id}.txt'
    shell: 'cp {input} {output}'
        
rule create_file:
    output: 'original_data/file{id}.txt'
    shell: 'touch {output}'
# -

# For running the workflow now, we need to specify which file we actually need as a final result and snakemake takes care of the individual steps to generate that file. We specify the desired output file as a snakemake argument:

# + language="sh"
# snakemake new_data/fileZ.txt --dag | dot | display
# snakemake new_data/fileZ.txt --dag | dot -Tsvg > dag2.svg
# snakemake new_data/fileZ.txt
# -

# To generate a set of output files, we can either request these individually when running snakemake, e.g. using `snakemake -np new_folder/file{0,1,2,3,4,5,6,7,8,9}.txt`. In case the workflow output is not being changed frequently, it is also possible to add a final rule (conventionally named 'all'), which requests all desired output files of the workflow:

# %%writefile Snakefile
rule all:
    input: expand('new_data/file{id}.txt', id=range(10))
rule copy_to_new_folder:
    input: 'original_data/file{id}.txt'
    output: 'new_data/file{id}.txt'
    shell: 'cp {input} {output}'
rule create_file:
    output: 'original_data/file{id}.txt'
    shell: 'touch {output}'

# + language="sh"
# snakemake --dag | dot | display
# snakemake --dag | dot -Tsvg > dag3.svg
# -

# Here I used the snakemake function `expand`, which extends a given statement (here `new_folder_file{id}.txt`) for all combinations of parameters provided (here `id` values from 0 to 10). This permits to easily applied a set of rules to a number of different files.

# ![DAG](dag3.svg)

# Typically, analysis is a bit more complicated than creating empty files and copying them from A to B using shell commands. Snakemake also support a number of different execution methods
# * in a shell (as used above)
# * in Python (using run:)
# * run Python/R/Markdown scripts directly (using script:)
# As an example we can use a small Python script to generate our initial data files and store a (randomly generated) value. The Python script would look like this:

# +
# %%writefile generate_data.py
import sys
import numpy as np

def generate_random_data(output_filename):
    # write a random number in an output file
    with open(output_filename, "w") as f:
        f.write(np.random.random())

# extracting the output filename from the command line parameters provided
output_filename = sys.argv[1]
generate_random_data(output_filename)
# -

# The corresponding snakemake rule now needs to provide the argument to the `generate_data.py` script:

# +
# %%writefile Snakefile

rule all:
    input: expand('new_data/file{id}.txt', id=range(10))
        
rule copy_to_new_folder:
    input: 'original_data/file{id}.txt'
    output: 'new_data/file{id}.txt'
    shell: 'cp {input} {output}'
        
rule generate_data:
    output: 'original_data/file{id}.txt'
    run: 'generate_data.py {output}'

# + language="sh"
# snakemake --dag | dot | display
# snakemake --dag | dot -Tsvg > dag4.svg
# -

# ![DAG](dag4.svg)

# Additional features worth having a look at
# * conda integration (--use-conda flag)
# * test runs (--dryrun)
# * print shell commands (-p)
# * includes of other Snakefiles
# * configuration files
# * rulegraph (--rulegraph)
# * remote files (e.g. Amazon Simple Storage Service, Google Cloud Storage, SSH, HTTP, FTP, Dopbox, WebDAV, ...)
# * [snakemake documentation](https://snakemake.readthedocs.io) and [FAQ](https://snakemake.readthedocs.io/en/stable/project_info/faq.html)

# ## Utilizing Snakemake for Data Analysis

# ### The electrophysiology analysis toolkit
# <img src=https://elephant.readthedocs.io/en/latest/_static/elephant_logo_sidebar.png alt="ElephantLogo" style="width: 300px;"/>  
# We will present a simple data analysis workflow based on a published dataset of complex electrophysiology data using the electrophysiology analysis toolkit [Elephant](http://neuralensemble.org/elephant). Elephant provides a number of methods for spiketrain and LFP analysis. Elephant contains a number of methods focussing on analysis of higher order correlation in spiking activity and the statistical evaluation of these. Elephant is based on the [Neo](https://neo.readthedocs.io) data framework, which we will first introduce based on artificially generated spiking data with a known correlation.

# ### Neo - the data framework <img src=http://neo.readthedocs.org/en/latest/_images/neologo.png alt="NeoLogo" style="width: 400px;"/>    

# + [markdown] slideshow={"slide_type": "skip"}
# ![NeoDataModel](https://neo.readthedocs.io/en/0.7.0/_images/base_schematic.png)
# *Neo* provides a set of classes for representing time series data (`SpikeTrain`, `AnalogSignalArray`, etc.), for representing the hierarchical arrangement of data in an experiment (`Segment`, `Block`) and for representing the relationship between spike trains and recording channels (`ChannelIndex`, `Unit`). *Neo* is used by a number of data visualization tools ([OpenElectrophy](http://neuralensemble.org/OpenElectrophy), [SpykeViewer](https://spyke-viewer.readthedocs.org/)) and by the [PyNN](http://neuralensemble.org/PyNN) metasimulator. *Neo* offers reading capabilities for a large number of proprietary electrophysiology data formats and conversion capability to generic open source data formats/models.
# <img src=images/IODiagram.svg alt="IODiagram" style="width: 600px;"/>
# -

# The *Neo* data model consists of container (`Block`, `Segment`, `ChannelIndex`, `Unit`) and data objects (`AnalogSignals`, `SpikeTrain`, `Epoch`, `Events`). Container objects provide the relation between the data stored in the *Neo* structure. Here, we consider only `Block` and `Segment` objects as containers. `Block`s are designed to contain everything related to a whole recording session and link to `Segment`s, which hold data belonging to a common time frame. For data objects we will use `SpikeTrain`s, designed to capture the time series describing the occurrence of spikes together with the corresponding waveforms. All *Neo* data objects are derivatives of NumPy arrays enhanced with the handling of physical quantities, minimal metadata as well as a generic mechanism to add custom metadata.
# <img src=https://neo.readthedocs.io/en/0.7.1/_downloads/simple_generated_diagram.svg alt="ClassDiagram" style="width: 800px;"/>
#

# Performing the analysis requires minimum 3 steps:
# * retrieving / generating the data
# * running the main analysis
# * generating result plots

# ### Generating ground-truth data
# Before running the analysis on experimental data, let's generate some data with known ground truth. Here Elephant provides a number of methods to generate spiketrain activity with defined statistical properties. The simplest would be to have independent spike time (Poisson process). Since we later want to detect higher order correlations, we will use a Poisson process as background activity and add correlated spikes to it.

# imports for data handling and visualization
from quantities import Hz, ms
from elephant.spike_train_generation import homogeneous_poisson_process, compound_poisson_process
import neo
import numpy as np
import matplotlib.pyplot as plt

# Neo and Elephant are handling physical units consistently during the analysis by using the Python module [quantities](https://python-quantities.readthedocs.io). This also requires parameters to be supplied in the correct dimension, such that the physical units can be matched during analysis. In general Neo objects capture all minimal information relevant for the interpretation of the data. In case of the spiketrain, this encompasses the start and stop times of recording/spike extraction as well the sampling rate and potential custom annotations in form of a dictionary.

spiketrain = homogeneous_poisson_process(20*Hz, 0*ms, 1000*ms)
print('The spiketrain', spiketrain)
print('Spiketrain attributes and physical units')
print(['{}: {}'.format(att, getattr(spiketrain, att))
       for att in ['t_start', 't_stop', 'sampling_rate', 'annotations']])
print(['{}: {}'.format(att, getattr(spiketrain, att).units.dimensionality)
       for att in ['t_start', 't_stop', 'sampling_rate']])

# An a first rule in our new workflow, let's generate multiple datasets with spiking activity and save them for future analysis steps. From the variety of file formats supported by Neo NIX has an hdf5 backend. Let's implement a virtual expiment, generating 100 Poisson spiketrains stored in the NIX framework:

# +
import neo
import quantities as pq
from elephant.spike_train_generation import homogeneous_poisson_process, compound_poisson_process
def generate_data(output_file, n=10):
    with neo.io.NixIO(output_file, 'ow') as io:
        # generate neo structure
        block = neo.Block(experiment='poisson')
        block.segments.append(neo.Segment(name='trial 1'))
        # generate correlated spike trains
        sts = compound_poisson_process(rate=5*pq.Hz, A=[0]+[0.90]+[0]*9+[0.1], t_stop=10*pq.s)
        # add background poisson spike trains
        for i in range(89):
            sts.append(homogeneous_poisson_process(rate=5*pq.Hz, t_stop=10*pq.s))
        block.segments[0].spiketrains.extend(sts)
        block.create_relationship()
        io.write_block(block)
    
generate_data('original_data/example_dataset.nix')
# -

# We export this piece of code into a standalone script so we can use it in the snakemake workflow:

# +
# %%writefile generate_poisson.py
import sys
import neo
import quantities as pq
from quantities import Hz, ms
from elephant.spike_train_generation import homogeneous_poisson_process, compound_poisson_process
def generate_data(output_file, n=10):
    with neo.io.NixIO(output_file, 'ow') as io:
        # generate neo structure
        block = neo.Block(experiment='poisson')
        block.segments.append(neo.Segment(name='trial 1'))
        # generate correlated spike trains
        sts = compound_poisson_process(rate=5*pq.Hz, A=[0]+[0.90]+[0]*9+[0.1], t_stop=10*pq.s)
        # add background poisson spike trains
        for i in range(89):
            sts.append(homogeneous_poisson_process(rate=5*pq.Hz, t_stop=10*pq.s))
        block.segments[0].spiketrains.extend(sts)
        block.create_relationship()
        io.write_block(block)
    
if __name__=='__main__':
    generate_data(*sys.argv[1:])

# + language="sh"
# python generate_poisson.py original_data/example_dataset.nix
# -

# The first rule in the workflow knows how to utilize the script to generate 10 new datasets. We limit the ids of the generated datasets from 0 to 9 to be able to use higher ids for different datasets.

# +
# %%writefile Snakefile
rule all:
    input: expand('original_data/dataset{id}.nix', id=range(10))
        
rule generate_data:
    output: 'original_data/dataset{id,[0-9]}.nix'
    shell: 'python generate_poisson.py {output}'
# -

# And running the workflow. The output files should appear [here](original_data). Have a look at the shell output of snakemake. It should list the number and type of jobs which are going to be run and the wildcards and status of each job.

# + language="sh"
# snakemake -pr

# + language="sh"
# snakemake --dag | dot | display
# snakemake --dag | dot -Tsvg > dag5.svg
# -

# ![DAG](dag5.svg)

# To visualize the generated data, we implement a utility script to load the data and a second script plot the data using Matplotlib.

# +
# %%writefile plot_data.py
import sys
import numpy as np
import neo
import matplotlib.pyplot as plt
def plot_data(data_filename, plot_filename):
    # opening the data file
    with neo.io.NixIO(data_filename, 'ro') as io:
        block = io.read_block()
    # plot spiketrains
    for i, spiketrain in enumerate(block.segments[0].spiketrains):
        plt.plot(spiketrain, [i]*len(spiketrain), 'C0.')
        # this is for later steps in the analysis: plot patterns if already detected
        if 'pattern' in spiketrain.annotations:
            pattern_spikes = spiketrain[np.where(spiketrain.annotations['pattern'])]
            plt.plot(pattern_spikes, [i]*len(pattern_spikes), 'r.')
    plt.xlabel('Time [{}]'.format(spiketrain[0].units.dimensionality.latex))
    plt.ylabel('Spiketrains')
    plt.savefig(plot_filename)
    
if __name__=='__main__':
    plot_data(*sys.argv[1:])
# -

# Now we can visualize the previously generated data. The introduced correlation should be visible already by eye in this arrangement of the spiketrains. In case of random reordering however this looks different.

from plot_data import plot_data
plot_data('original_data/dataset2.nix', 'original_data/dataset2.png')

# For the context of this analysis we teach snakemake how to go from arbitraty datasets in the nix format to a plot of the data. Now the exact data location of the nix file is not relevant any more since we introduced an additional `folder` wildcard.

# +
# %%writefile -a Snakefile

rule plot_data:
    input: '{folder}/dataset{id}.nix'
    output: '{folder}/dataset{id}.png'
    shell: 'python plot_data.py {input} {output}'

# + language="sh"
# snakemake original_data/dataset0.png 
# -

# ### Running SPADE
# In the next step we want to run an analysis on the data, which is detecting the injected higher order correlation in the data. One of the methods suited for the detection of spatio-temporal patterns is `SPADE` [(Quaglio et al, 2017)](https://doi.org/10.3389/fncom.2017.00041)
# .

# +
# %%writefile run_spade.py
import sys
import numpy as np
import neo
import quantities as pq
import elephant.spade


def clean_nix_annotations(block):
    '''helper function needed for neo 0.7.1'''
    # removing unnessesary annotations to permit block to be saved again via nix
    objs = [block] + block.filter(data=True, container=True)
    for obj in objs:
        if 'nix_name' in obj.annotations:
            obj.annotations.pop('nix_name')
        if 'neo_name' in obj.annotations:
            obj.annotations.pop('neo_name')
    

def run_spade(data_filename, output_filename):
    # loading dataset
    with neo.io.get_io(data_filename, 'rw') as io:
        block = io.read_block()
    spiketrains = block.segments[0].spiketrains
    
    # run spade analysis
    patterns = elephant.spade.spade(
        data=spiketrains, binsize=1*pq.ms, winlen=1, dither=5*pq.ms,
        min_spikes=3, n_surr=10, psr_param=[0,0,3],
        output_format='patterns')['patterns']
    
    # convert first pattern into annotations of spiketrains
    # default: no spike belongs to any pattern
    for spiketrain in spiketrains:
        spiketrain.annotate(pattern=np.array([False]*len(spiketrain)))
        
    # convert first detected pattern to annotations of spiketrains
    for pattern_id, pattern in enumerate(patterns):
        neurons = pattern['neurons']
        times = pattern['times']
        lags = pattern['lags']
        for n, neuron in enumerate(neurons):
            if n == 0:
                pattern_times = times
            else:
                pattern_times = times + lags[n-1]
            for pattern_time in pattern_times:
                spike_idx = np.abs(spiketrains[neuron]-pattern_time).argmin()
                spiketrains[neuron].annotations['pattern'][spike_idx] = True

    # overwrite original neo file
    clean_nix_annotations(block)
    with neo.io.NixIO(output_filename, 'ow') as io2: 
        io2.write_block(block)
        
if __name__=='__main__':
    run_spade(*sys.argv[1:])
# -

# We add a rule to the snakemake workflow for runnig spade:

# +
# %%writefile -a Snakefile

rule run_spade:
    input: 'original_data/dataset{id}.nix'
    output: 'results/dataset{id}.nix'
    shell: 'python run_spade.py {input} {output}'
# -

# Now we can run spade by simply requesting the analysis result from snakemake

# + language="sh"
# snakemake results/dataset0.nix
# -

# ## Apply analysis to experimental data <img src=https://web.gin.g-node.org/img/favicon.png alt="GIN" style="width: 100px;"/>
# In the next step we want to apply the same workflow to a set of experimental data. The data are publicly available at the data hosting platform [GIN](https://web.gin.g-node.org/) and are described in detail in [Brochier (2018)](https://doi.org/10.1038/sdata.2018.55). We are going to download two specific files of the dataset, which are saved in BlackRock nev/ns2 format. Since Neo also supports BlackRock formats, we use it to extract the relevant data and save them in the nix format to be consistent with the simulated data.

# +
# %%writefile convert_to_nix.py
import sys
import neo
import quantities as pq

def convert_data_to_nix(input_filename, nix_filename):
    print(input_filename)
    '''extracting first seconds of first spiketrains into separate nix file''' 
    io = neo.io.BlackrockIO(input_filename)
    original_block = io.read_block()
    
    # Extracting spiketrain and append to new block
    block = neo.Block(experiment='monkey')
    block.segments.append(neo.Segment(name='first seconds of recording'))
    new_spiketrains = [st.time_slice(0*pq.s,10*pq.s) for st in original_block.segments[-1].spiketrains[:100]]
    block.segments[0].spiketrains = new_spiketrains
    
    # save new block
    with neo.io.NixIO(nix_filename, 'ow') as io2:
        io2.write_block(block)
        
if __name__=='__main__':
    convert_data_to_nix(*sys.argv[1:])
# -

# For fetching the data we directly download the data using `wget`. A more elegant version would involve installation of `git annex` or the `gin` client to be consistent with the version management handled by gin. We also add a generic rule which permits the conversion of BlackRock files to nix files while cropping only a selection of the data. The final rule `run_complete` requests visualizations of all analysis runs and therefore triggers the whole workflow.

# +
# %%writefile -a Snakefile

rule download_data:
    output:
        nev = 'original_data/dataset10.nev',
        ns2 = 'original_data/dataset10.ns2'
    shell: '''
    wget -O {output.nev} https://web.gin.g-node.org/INT/multielectrode_grasp/raw/24cd5caee3ae79066ca37844cab931d04dcad977/datasets/i140703-001-03.nev
    wget -O {output.ns2} https://web.gin.g-node.org/INT/multielectrode_grasp/raw/24cd5caee3ae79066ca37844cab931d04dcad977/datasets/i140703-001.ns2
    '''
        
rule convert_to_nix:
    input: '{filename}.nev',
           '{filename}.ns2'
    output: '{filename}.nix'
    shell: 'python convert_to_nix.py {wildcards.filename} {output}'
        
rule run_complete:
    input: expand('results/dataset{id}.png', id=range(11))

# + language="sh"
# snakemake run_complete
# -

# The whole workflow encompasses now artificial data generation as well as retrieval of experimental data as well as a common analysis and visualization.

# + language="sh"
# snakemake run_complete --dag | dot | display
# snakemake run_complete --dag | dot -Tsvg > run_complete.svg
# -

# ![DAG](run_complete.svg)



StarcoderdataJupyterScriptsDedupFiltered