from elephant.spike_train_generation import             homogeneous_poisson_process, homogeneous_gamma_process
import quantities as pq
import matplotlib.pyplot as plt
from viziphant.rasterplot import rasterplot_rates

spiketrains1 = [homogeneous_poisson_process(rate=10 * pq.Hz)
                for _ in range(100)]
spiketrains2 = [homogeneous_gamma_process(a=3, b=10 * pq.Hz)
                for _ in range(100)]
for i, (st1, st2) in enumerate(zip(spiketrains1, spiketrains2)):
    if i % 2 == 1:
        st1.annotations['parity'] = 'odd'
        st2.annotations['parity'] = 'odd'
    else:
        st1.annotations['parity'] = 'even'
        st2.annotations['parity'] = 'even'

# plot separates the lists and the annotation values within each list
rasterplot_rates([spiketrains1, spiketrains2], key_list=['parity'],
                  groupingdepth=2, labelkey='0+1')