Merge 570da13 into 76ab71a

brian2/codegen/runtime/cython_rt/templates/synapses_create_generator.pyx

@@ -38,10 +38,19 @@ cdef void _flush_buffer(buf, dynarr, int buf_len):
     cdef size_t newsize
 
     # The following variables are only used for probabilistic connections
+    {% if iterator_func=='sample' %}
+    cdef int _iter_sign
+    {% if iterator_kwds['sample_size'] == 'fixed' %}
+    cdef int _n_selected
+    cdef int _n_dealt_with
+    cdef int _n_total
+    cdef double _U
+    {% else %}
     cdef bool _jump_algo
     cdef double _log1p, _pconst
     cdef size_t _jump
-
+    {% endif %}
+    {% endif %}
     # scalar code
     _vectorisation_idx = 1
     {{scalar_code['setup_iterator']|autoindent}}
@@ -61,25 +70,47 @@ cdef void _flush_buffer(buf, dynarr, int buf_len):
         {% if iterator_func=='range' %}
         for {{iteration_variable}} in range(_iter_low, _iter_high, _iter_step):
         {% elif iterator_func=='sample' %}
+        {% if iterator_kwds['sample_size'] == 'fixed' %}
+        # Selection sampling technique
+        # See section 3.4.2 of Donald E. Knuth, AOCP, Vol 2, Seminumerical Algorithms
+        _n_selected = 0
+        _n_dealt_with = 0
+        if _iter_step > 0:
+            _n_total = (_iter_high - _iter_low - 1) // _iter_step + 1
+        else:
+            _n_total = (_iter_low - _iter_high - 1) // -_iter_step + 1
+        for {{iteration_variable}} in range(_iter_low, _iter_high, _iter_step):
+            _U = _rand(_vectorisation_idx)
+            if (_n_total - _n_dealt_with)*_U >= _iter_size - _n_selected:
+                _n_dealt_with += 1
+                continue
+            _n_selected += 1
+            _n_dealt_with += 1
+        {% else %}
         if _iter_p==0:
             continue
+        if _iter_step < 0:
+            _iter_sign = -1
+        else:
+            _iter_sign = 1
         _jump_algo = _iter_p<0.25
         if _jump_algo:
             _log1p = log(1-_iter_p)
         else:
             _log1p = 1.0 # will be ignored
         _pconst = 1.0/_log1p
-        {{iteration_variable}} = _iter_low-1
-        while {{iteration_variable}}+1<_iter_high:
-            {{iteration_variable}} += 1
+        {{iteration_variable}} = _iter_low-_iter_step
+        while _iter_sign*({{iteration_variable}} + _iter_step) < _iter_sign*_iter_high:
+            {{iteration_variable}} += _iter_step
             if _jump_algo:
-                _jump = int(floor(log(_rand(_vectorisation_idx))*_pconst))*_iter_step
+                _jump = <int>(log(_rand(_vectorisation_idx))*_pconst)*_iter_step
                 {{iteration_variable}} += _jump
-                if {{iteration_variable}}>=_iter_high:
+                if _iter_sign*{{iteration_variable}} >= _iter_sign*_iter_high:
                     break
             else:
                 if _rand(_vectorisation_idx)>=_iter_p:
                     continue
+        {% endif %}
         {% endif %}
 
             {{vector_code['create_j']|autoindent}}

brian2/codegen/runtime/numpy_rt/templates/synapses_create_generator.py_

@@ -21,7 +21,7 @@ except:
 
     from numpy.random import binomial as _binomial
     import bisect as _bisect
-    def _sample_without_replacement(low, high, step, p):
+    def _sample_without_replacement(low, high, step, p=None, size=None):
         if p == 0:
             return _numpy.array([], dtype=_numpy.int32)
         if p == 1:
@@ -30,8 +30,17 @@ except:
             n = (high-low-1)//step+1
         else:
             n = (low-high-1)//-step+1
-        if n <= 0:
+        if n <= 0 or (size is not None and size <= 0):
             return _numpy.array([], dtype=_numpy.int32)
+        if n == size:
+            return _numpy.arange(low, high, step)
+        if p is None:
+            if step == 0:
+                _choice = _numpy.random.choice(n,  size=size, replace=False) + low
+            else:
+                _chose_from = _numpy.arange(low, high, step)
+                _choice = _numpy.random.choice(_chose_from, size=size, replace=False)
+            return _numpy.sort(_choice)
         if n <= 1000 or p > 0.25: # based on rough profiling
             samples, = (_np_rand(n)<p).nonzero()
         else:
@@ -89,8 +98,12 @@ for _i in range(N_pre):
     {% if iterator_func=='range' %}
     {{iteration_variable}} = _numpy.arange(_iter_low, _iter_high, _iter_step)
     {% elif iterator_func=='sample' %}
+    {% if iterator_kwds['sample_size'] == 'fixed' %}
+    {{iteration_variable}} = _sample_without_replacement(_iter_low, _iter_high, _iter_step, None, size=_iter_size)
+    {% else %}
     {{iteration_variable}} = _sample_without_replacement(_iter_low, _iter_high, _iter_step, _iter_p)
     {% endif %}
+    {% endif %}
 
     _vectorisation_idx = {{iteration_variable}}
     {{vector_code['create_j']|autoindent}}

brian2/devices/cpp_standalone/templates/synapses_create_generator.cpp

@@ -53,6 +53,7 @@
         long _uiter_high;
         long _uiter_step;
         {% if iterator_func=='sample' %}
+        long _uiter_size;
         double _uiter_p;
         {% endif %}
         {
@@ -61,32 +62,59 @@
             _uiter_high = _iter_high;
             _uiter_step = _iter_step;
             {% if iterator_func=='sample' %}
+            {% if iterator_kwds['sample_size'] == 'fixed' %}
+            _uiter_size = _iter_size;
+            {% else %}
             _uiter_p = _iter_p;
             {% endif %}
+            {% endif %}
         }
         {% if iterator_func=='range' %}
         for(long {{iteration_variable}}=_uiter_low; {{iteration_variable}}<_uiter_high; {{iteration_variable}}+=_uiter_step)
         {
         {% elif iterator_func=='sample' %}
+        {% if iterator_kwds['sample_size'] == 'fixed' %}
+        // Selection sampling technique
+        // See section 3.4.2 of Donald E. Knuth, AOCP, Vol 2, Seminumerical Algorithms
+        int _n_selected = 0;
+        int _n_dealt_with = 0;
+        int _n_total;
+        if (_uiter_step > 0)
+            _n_total = (_uiter_high - _uiter_low - 1) / _uiter_step + 1;
+        else
+            _n_total = (_uiter_low - _uiter_high - 1) / -_uiter_step + 1;
+        for(long {{iteration_variable}}=_uiter_low; {{iteration_variable}}<_uiter_high; {{iteration_variable}}+=_uiter_step)
+        {
+            const double _U = _rand(_vectorisation_idx);
+            if ((_n_total - _n_dealt_with)*_U >= _uiter_size - _n_selected)
+            {
+                _n_dealt_with += 1;
+                continue;
+            }
+            _n_selected += 1;
+            _n_dealt_with += 1;
+        {% else %}
         if(_uiter_p==0) continue;
+        const int _iter_sign = _uiter_step > 0 ? 1 : -1;
         const bool _jump_algo = _uiter_p<0.25;
         double _log1p;
         if(_jump_algo)
             _log1p = log(1-_uiter_p);
         else
             _log1p = 1.0; // will be ignored
         const double _pconst = 1.0/log(1-_uiter_p);
-        for(long {{iteration_variable}}=_uiter_low; {{iteration_variable}}<_uiter_high; {{iteration_variable}}++)
+        for(long {{iteration_variable}}=_uiter_low; _iter_sign*{{iteration_variable}}<_iter_sign*_uiter_high; {{iteration_variable}} += _uiter_step)
         {
             if(_jump_algo) {
                 const double _r = _rand(_vectorisation_idx);
                 if(_r==0.0) break;
                 const int _jump = floor(log(_r)*_pconst)*_uiter_step;
                 {{iteration_variable}} += _jump;
-                if({{iteration_variable}}>=_uiter_high) continue;
+                if (_iter_sign*{{iteration_variable}} >= _iter_sign * _uiter_high) continue;
             } else {
-                if(_rand(_vectorisation_idx)>=_uiter_p) continue;
+                if (_rand(_vectorisation_idx)>=_uiter_p) continue;
             }
+        {% endif %}
         {% endif %}
             long __j, _j, _pre_idx, __pre_idx;
             {

brian2/synapses/synapses.py

@@ -1696,10 +1696,6 @@ def _add_synapses_generator(self, j, n, skip_if_invalid=False, namespace=None, l
         template_kwds.update(parsed)
         template_kwds['skip_if_invalid'] = skip_if_invalid
 
-        if (parsed['iterator_func'] == 'sample' and
-                    parsed['iterator_kwds']['sample_size']=='fixed'):
-            raise NotImplementedError("Fixed sample size not implemented yet.")
-
         abstract_code = {'setup_iterator': '',
                          'create_j': '',
                          'create_cond': '',