Merge pull request #4 from tihbe/parity_task

Added generated bit parity task

Author: tihbe

README.md

@@ -45,7 +45,7 @@ You can install the latest version of this package with:
 pip install git+https://github.com/tihbe/python-ebdataset.git
 ```
 
-# Usage
+# Getting started
 
 In the code:
 ```python

ebdataset/audio/__init__.py

@@ -1 +1,2 @@
+"""This subpackage regroups audio-based spiking dataset"""
 from .ntidigits import NTidigits

ebdataset/generated/__init__.py

@@ -0,0 +1,2 @@
+"""This subpackage regroups self-generated spiking dataset"""
+from .parity import ParityTask

ebdataset/generated/parity.py

@@ -0,0 +1,70 @@
+import numpy as np
+import torch.utils.data as data
+from quantities import hertz, ms, second
+
+
+class ParityTask(data.IterableDataset):
+    """Create a spike-based population encoded parity (or n-bits xor) task of max_iter samples
+    if max_iter is not specified or is np.inf, this dataset will keep generating samples forever
+    each HIGH bits is encoded with high_freq poisson-sampled spikes of shape features_per_bit x duration_per_bit
+    LOW bits and background-noise is encoded with low_freq poisson-sampled spikes for the remaining of the sample_duration
+    bits are encoded both temporally and spatially
+    """
+
+    def __init__(
+        self,
+        seed=0x1B,
+        low_freq=2 * hertz,
+        high_freq=20 * hertz,
+        sample_duration=2 * second,
+        number_of_bits=2,
+        features_per_bit=50,
+        duration_per_bit=0.5 * second,
+        dt=1 * ms,
+        max_iter=np.inf,
+        as_recarray=True,
+    ):
+        self.seed = seed
+        self.max_iter = max_iter
+        self.gen = np.random.RandomState(seed=self.seed)
+        self.low_freq = float((low_freq * dt).simplified)
+        self.high_freq = float((high_freq * dt).simplified)
+        self.sample_duration = int((sample_duration / dt).simplified)
+        self.duration_per_bit = int((duration_per_bit / dt).simplified)
+        self.number_of_bits = number_of_bits
+        self.features_per_bit = features_per_bit
+        self.as_recarray = as_recarray
+
+    def __iter__(self):
+        worker_info = data.get_worker_info()
+        m = 1
+        if worker_info is not None:  # multi-process data loading, re-seed the iterator
+            self.gen = np.random.RandomState(seed=worker_info.id + self.seed)
+            m = worker_info.num_workers
+
+        i = 0
+        while i < self.max_iter / m:
+            i += 1
+
+            bits = self.gen.randint(0, 2, size=self.number_of_bits)
+            y = np.sum(bits) % 2
+
+            spike_train = self.gen.poisson(
+                lam=self.low_freq, size=(self.number_of_bits * self.features_per_bit, self.sample_duration)
+            )
+
+            for b in range(self.number_of_bits):
+                if bits[b]:
+                    spike_train[
+                        b * self.features_per_bit : (b + 1) * self.features_per_bit,
+                        b * self.duration_per_bit : (b + 1) * self.duration_per_bit,
+                    ] = self.gen.poisson(lam=self.high_freq, size=(self.features_per_bit, self.duration_per_bit))
+
+            if self.as_recarray:
+                addr, ts = np.nonzero(spike_train)
+                sample = np.recarray(shape=len(ts), dtype=[("addr", addr.dtype), ("ts", ts.dtype)])
+                sample.addr = addr
+                sample.ts = ts
+                yield sample, y
+            else:
+                yield spike_train, y

ebdataset/vision/__init__.py

@@ -1,3 +1,4 @@
+"""This subpackage regroups vision-based spiking dataset"""
 from .ibm_gesture import IBMGesture, H5IBMGesture
 from .ini_ucf50 import INIUCF50
 from .ncaltech101 import NCaltech101