Changed units system from quantities to pint

Author: Ismael Balafrej

README.md

@@ -51,7 +51,7 @@ In the code:
 ```python
 from ebdataset.vision import NMnist
 from ebdataset.vision.transforms import ToDense
-from quantities import ms
+from ebdataset import ms
 
 # With sparse representation:
 for spike_train, label in NMnist(path):

ebdataset/__init__.py

@@ -0,0 +1,4 @@
+import ebdataset.vision as vision
+import ebdataset.audio as audio
+import ebdataset.generated as generated
+from .utils.units import *

ebdataset/generated/parity.py

@@ -1,6 +1,6 @@
 import numpy as np
 import torch.utils.data as data
-from quantities import hertz, ms, second
+from ..utils.units import hertz, ms, second, wunits
 
 
 class ParityTask(data.IterableDataset):
@@ -11,6 +11,7 @@ class ParityTask(data.IterableDataset):
     bits are encoded both temporally and spatially if sequential=True, otherwise only spatially
     """
 
+    @wunits(None, (None, None, hertz, hertz, second, None, None, second, second, None, None, None))
     def __init__(
         self,
         seed=0x1B,
@@ -27,35 +28,35 @@ def __init__(
     ):
         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.rand = np.random.RandomState(seed=self.seed)
+        self.low_freq = low_freq * dt
+        self.high_freq = high_freq * dt
+        self.sample_duration = int(sample_duration / dt)
+        self.duration_per_bit = int(duration_per_bit / dt)
         self.number_of_bits = number_of_bits
         self.features_per_bit = features_per_bit
         self.as_recarray = as_recarray
         self.sequential = sequential
         if sequential:
             assert (
-                duration_per_bit * number_of_bits >= sample_duration
+                duration_per_bit * number_of_bits <= sample_duration
             ), "Sample duration is not enough to contain every bits"
 
     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)
+            self.rand = 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)
+            bits = self.rand.randint(0, 2, size=self.number_of_bits)
             y = np.sum(bits) % 2
 
-            spike_train = self.gen.poisson(
+            spike_train = self.rand.poisson(
                 lam=self.low_freq, size=(self.number_of_bits * self.features_per_bit, self.sample_duration)
             )
 
@@ -68,7 +69,7 @@ def __iter__(self):
                     )
                     spike_train[
                         b * self.features_per_bit : (b + 1) * self.features_per_bit, time_pos
-                    ] = self.gen.poisson(lam=self.high_freq, size=(self.features_per_bit, self.duration_per_bit))
+                    ] = self.rand.poisson(lam=self.high_freq, size=(self.features_per_bit, self.duration_per_bit))
 
             if self.as_recarray:
                 ts, addr = np.nonzero(spike_train.T)

ebdataset/utils/units.py

@@ -0,0 +1,15 @@
+from pint import UnitRegistry as _UR
+
+_reg = _UR()
+
+wunits = _reg.wraps
+
+## Global time management namespace
+second = s = _reg.s
+millisecond = ms = _reg.ms
+microsecond = us = _reg.us
+nanosecond = ns = _reg.ns
+killosecond = ks = _reg.ks
+hertz = Hz = _reg.Hz
+millihertz = mhertz = mHz = _reg.mHz
+kilohertz = khertz = kHz = _reg.kHz

ebdataset/vision/ini_roshambo.py

@@ -5,7 +5,7 @@
 from .parsers.aedat import readAEDATv2_davies
 from torch.utils.data.dataset import Dataset
 from .type import DVSSpikeTrain
-from quantities import us
+from ..utils.units import us, wunits
 
 
 class INIRoshambo(Dataset):
@@ -66,11 +66,12 @@ def convert(self, out_path, verbose=False):
 
         return INIRoshambo(out_path, with_backgrounds=self.with_backgrounds, transforms=self.transforms)
 
+    @wunits(None, (None, None, us, None))
     def split_to_subsamples(self, out_path, duration_per_sample, verbose=False):
         if not (".h5" in out_path):
             out_path += ".h5"
 
-        duration_per_sample = int(duration_per_sample.rescale(us).magnitude)
+        duration_per_sample = int(duration_per_sample)
 
         with File(out_path, "w-", libver="latest") as f_hndl:
             for i, (sample, label) in enumerate(tqdm(self, disable=not verbose)):

ebdataset/vision/transforms.py

@@ -3,7 +3,7 @@
 """
 import torch
 import numpy as np
-import quantities as units
+from ..utils.units import second, us, ms, wunits
 from torchvision.transforms import Compose
 
 
@@ -32,8 +32,9 @@ def __call__(self, sparse_spike_train):
 class MaxTime(object):
     """Limit the time of a 2d sparse spike train"""
 
-    def __init__(self, max_time: units.UnitTime, dt: units.UnitTime = 1 * units.us):
-        self.max = (max_time.rescale(dt.units) / dt).magnitude
+    @wunits(None, (None, second, second))
+    def __init__(self, max_time, dt=1 * us):
+        self.max = max_time / dt
 
     def __call__(self, sparse_spike_train):
         mask = sparse_spike_train.ts < self.max
@@ -50,14 +51,15 @@ class ToDense(object):
     """Transform a sparse spike train to a dense torch tensor of shape (x, y, p, time)
     with time unit defined by dt. Time accumulation is done with a max function."""
 
+    @wunits(None, (None, second))
     def __init__(
         self,
-        dt: units.UnitTime,  # Time scale of dense tensor
+        dt,  # Time scale of dense tensor
     ):
         self.dt = dt
 
     def __call__(self, sparse_spike_train):
-        time_scale = ((sparse_spike_train.time_scale * units.second).rescale(self.dt.units) / self.dt).magnitude
+        time_scale = sparse_spike_train.time_scale / self.dt
         duration = np.ceil(sparse_spike_train.duration * time_scale).astype(int)
         dense_spike_train = torch.zeros((sparse_spike_train.width, sparse_spike_train.height, 2, duration))
 

setup.py

@@ -7,7 +7,7 @@
 
 setup(
     name="ebdataset",
-    version="0.0.2",
+    version="0.1.0",
     author="Ismael Balafrej - NECOTIS",
     author_email="[email protected]",
     description="An event based dataset loader under one common API.",
@@ -17,7 +17,7 @@
     packages=find_packages(),
     install_requires=[
         "numpy>=1.14.3",
-        "quantities>=0.12.4",
+        "pint>=0.17",
         "tqdm>=4.45.0",
         "torch>=1.4.0",
         "torchvision>=0.5.0",