Merge 5966ed1 into 1b780ee

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

@@ -1,23 +1,29 @@
-# USES_VARIABLES { _source_neurons, _target_neurons, _synaptic_pre, _synaptic_post, _pre_synaptic, _post_synaptic }
+# USES_VARIABLES { _source_neurons, _target_neurons, _synaptic_pre, _synaptic_post,
+#                  _pre_synaptic, _post_synaptic, _source_offset, _target_offset }
 
 import numpy as np
 
+numpy_False = np.bool_(False)
+numpy_True = np.bool_(True)
+
 for i in range(_num_source_neurons):
     j = np.arange(_num_target_neurons)
     _vectorisation_idx = j
+    _presynaptic_idx = i + _source_offset
+    _postsynaptic_idx = j + _target_offset
     {% for line in code_lines %}
     {{line}}
     {% endfor %}
-
-    if _cond is False:
+    if _cond is False or _cond is numpy_False:
         continue
-    elif _cond is True:
-        _cond_nonzero = j
-    else:
-        _cond_nonzero = _cond.nonzero()[0]
 
     if not np.isscalar(_p) or _p != 1:
-        _cond_nonzero = _cond_nonzero[np.random.rand(len(_cond_nonzero)) < _p]
+        _cond_nonzero, = np.logical_and(_cond,
+                                       np.random.rand(len(_vectorisation_idx)) < _p).nonzero()
+    elif _cond is True or _cond is numpy_True:
+        _cond_nonzero = j
+    else:
+        _cond_nonzero, = _cond.nonzero()
 
     if not np.isscalar(_n):
         # The "n" expression involved j

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

@@ -0,0 +1,9 @@
+# USES_VARIABLES { _synaptic_pre, _synaptic_post }
+# ITERATE_ALL { _idx }
+
+_vectorisation_idx = _idx
+_presynaptic_idx = _synaptic_pre[:]
+_postsynaptic_idx = _synaptic_post[:]
+{% for line in code_lines %}
+{{line}}
+{% endfor %}

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

@@ -0,0 +1,11 @@
+# USES_VARIABLES {_synaptic_pre, _synaptic_post}
+# ITERATE_ALL { _idx }
+
+_vectorisation_idx = _idx
+_presynaptic_idx = _synaptic_pre[:]
+_postsynaptic_idx = _synaptic_post[:]
+
+{% for line in code_lines %}
+{{line}}
+{% endfor %}
+_return_values, = _cond.nonzero()

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

@@ -0,0 +1,9 @@
+# USES_VARIABLES { _group_idx, _synaptic_pre, _synaptic_post,
+#                  _source_offset, _target_offset }
+_idx = _group_idx
+_vectorisation_idx = _idx
+_presynaptic_idx = _synaptic_pre[_idx] + _source_offset
+_postsynaptic_idx = _synaptic_post[_idx] + _target_offset
+{% for line in code_lines %}
+{{line}}
+{% endfor %}

brian2/codegen/runtime/weave_rt/templates/synapses_create.cpp

@@ -1,7 +1,7 @@
 {% macro main() %}
     // USES_VARIABLES { _synaptic_pre, _synaptic_post, _post_synaptic,
     //                  _pre_synaptic, _source_neurons, _target_neurons,
-    //                  rand}
+    //                  rand, _source_offset, _target_offset}
 
     //// SUPPORT CODE //////////////////////////////////////////////////////////
 	{% for line in support_code_lines %}
@@ -36,6 +36,8 @@
 		for(int j=0; j<_num_target_neurons; j++)
 		{
 		    const int _vectorisation_idx = j;
+		    const int _presynaptic_idx = i + _source_offset;
+		    const int _postsynaptic_idx = j + _target_offset;
 			// Define the condition
 			{% for line in code_lines %}
 			{{line}}

brian2/codegen/runtime/weave_rt/templates/synaptic_stateupdate.cpp

@@ -0,0 +1,46 @@
+////////////////////////////////////////////////////////////////////////////
+//// MAIN CODE /////////////////////////////////////////////////////////////
+
+{% macro main() %}
+    // USES_VARIABLES { _synaptic_pre, _synaptic_post }
+
+    ////// SUPPORT CODE ///
+	{% for line in support_code_lines %}
+	//{{line}}
+	{% endfor %}
+
+	////// HANDLE DENORMALS ///
+	{% for line in denormals_code_lines %}
+	{{line}}
+	{% endfor %}
+
+	////// HASH DEFINES ///////
+	{% for line in hashdefine_lines %}
+	{{line}}
+	{% endfor %}
+
+	///// POINTERS ////////////
+	{% for line in pointers_lines %}
+	{{line}}
+	{% endfor %}
+
+	//// MAIN CODE ////////////
+	for(int _idx=0; _idx<_num_idx; _idx++)
+	{
+	    const int _vectorisation_idx = _idx;
+	    const int _presynaptic_idx = _synaptic_pre[_idx];
+	    const int _postsynaptic_idx = _synaptic_post[_idx];
+		{% for line in code_lines %}
+		{{line}}
+		{% endfor %}
+	}
+{% endmacro %}
+
+////////////////////////////////////////////////////////////////////////////
+//// SUPPORT CODE //////////////////////////////////////////////////////////
+
+{% macro support_code() %}
+	{% for line in support_code_lines %}
+	{{line}}
+	{% endfor %}
+{% endmacro %}

brian2/codegen/runtime/weave_rt/templates/synaptic_variable_indexing.cpp

@@ -0,0 +1,54 @@
+////////////////////////////////////////////////////////////////////////////
+//// MAIN CODE /////////////////////////////////////////////////////////////
+
+{% macro main() %}
+    // USES_VARIABLES { _synaptic_pre, _synaptic_post }
+	////// SUPPORT CODE ///////
+	{% for line in support_code_lines %}
+	// {{line}}
+	{% endfor %}
+
+	////// HANDLE DENORMALS ///
+	{% for line in denormals_code_lines %}
+	{{line}}
+	{% endfor %}
+
+	////// HASH DEFINES ///////
+	{% for line in hashdefine_lines %}
+	{{line}}
+	{% endfor %}
+
+	///// POINTERS ////////////
+	{% for line in pointers_lines %}
+	{{line}}
+	{% endfor %}
+
+	//// MAIN CODE ////////////
+	int _cpp_numelements = 0;
+	// Container for all the potential indices
+	npy_int *_elements = (npy_int *)malloc(sizeof(npy_int) * _num_idx);
+	for(int _idx=0; _idx<_num_idx; _idx++)
+	{
+	    const int _vectorisation_idx = _idx;
+	    const int _presynaptic_idx = _synaptic_pre[_idx];
+	    const int _postsynaptic_idx = _synaptic_post[_idx];
+		{% for line in code_lines %}
+		{{line}}
+		{% endfor %}
+		if(_cond) {
+			_elements[_cpp_numelements++] = _idx;
+		}
+	}
+	npy_intp _dims[] = {_cpp_numelements};
+	PyObject *_numpy_elements = PyArray_SimpleNewFromData(1, _dims, NPY_INT, _elements);
+	return_val = _numpy_elements;
+{% endmacro %}
+
+////////////////////////////////////////////////////////////////////////////
+//// SUPPORT CODE //////////////////////////////////////////////////////////
+
+{% macro support_code() %}
+	{% for line in support_code_lines %}
+	// {{line}}
+	{% endfor %}
+{% endmacro %}

brian2/codegen/runtime/weave_rt/templates/synaptic_variable_set.cpp

@@ -0,0 +1,42 @@
+////////////////////////////////////////////////////////////////////////////
+//// MAIN CODE /////////////////////////////////////////////////////////////
+
+{% macro main() %}
+	// USES_VARIABLES { _group_idx, _synaptic_pre, _synaptic_post,
+    //                  _source_offset, _target_offset }
+	////// HANDLE DENORMALS ///
+	{% for line in denormals_code_lines %}
+	{{line}}
+	{% endfor %}
+
+	////// HASH DEFINES ///////
+	{% for line in hashdefine_lines %}
+	{{line}}
+	{% endfor %}
+
+	///// POINTERS ////////////
+	{% for line in pointers_lines %}
+	{{line}}
+	{% endfor %}
+
+	//// MAIN CODE ////////////
+	for(int _idx_group_idx=0; _idx_group_idx<_num_group_idx; _idx_group_idx++)
+	{
+		const int _idx = _group_idx[_idx_group_idx];
+		const int _presynaptic_idx = _synaptic_pre[_idx] + _source_offset;
+		const int _postsynaptic_idx = _synaptic_post[_idx] + _target_offset;
+		const int _vectorisation_idx = _idx;
+		{% for line in code_lines %}
+		{{line}}
+		{% endfor %}
+	}
+{% endmacro %}
+
+////////////////////////////////////////////////////////////////////////////
+//// SUPPORT CODE //////////////////////////////////////////////////////////
+
+{% macro support_code() %}
+	{% for line in support_code_lines %}
+	{{line}}
+	{% endfor %}
+{% endmacro %}

brian2/core/variables.py

@@ -222,11 +222,33 @@ def __repr__(self):
 
 
 class VariableView(object):
+    '''
+    A view on a variable that allows to treat it as an numpy array while
+    allowing special indexing (e.g. with strings) in the context of a `Group`.
+
+    Parameters
+    ----------
+    name : str
+        The name of the variable
+    variable : `Variable`
+        The variable description.
+    group : `Group`
+        The group to which this variable belongs
+    template : str
+        The template to use when setting variables with a string expression.
+    unit : `Unit`, optional
+        The unit to be used for the variable, should be `None` when a variable
+         is accessed without units (e.g. when stating `G.var_`).
+    level : int, optional
+        How much farther to go down in the stack to find the namespace.
+    '''
 
-    def __init__(self, name, variable, group, unit=None, level=0):
+    def __init__(self, name, variable, group, template,
+                 unit=None, level=0):
         self.name = name
         self.variable = variable
         self.group = group
+        self.template = template
         self.unit = unit
         self.level = level
 
@@ -254,7 +276,8 @@ def __setitem__(self, i, value):
         if isinstance(value, basestring):
             check_units = self.unit is not None
             self.group._set_with_code(variable, indices, value,
-                                      check_units, level=self.level + 1)
+                                      template=self.template,
+                                      check_units=check_units, level=self.level + 1)
         else:
             if not self.unit is None:
                 fail_for_dimension_mismatch(value, self.unit)
@@ -372,10 +395,16 @@ def set_value(self, value):
         self.value[:] = value
 
     def get_addressable_value(self, group, level=0):
-        return VariableView(self.name, self, group, None, level)
+        template = getattr(group, '_set_with_code_template',
+                           'group_variable_set')
+        return VariableView(self.name, self, group, template=template,
+                            unit=None, level=level)
 
     def get_addressable_value_with_unit(self, group, level=0):
-        return VariableView(self.name, self, group, self.unit, level)
+        template = getattr(group, '_set_with_code_template',
+                           'group_variable_set')
+        return VariableView(self.name, self, group, template=template,
+                            unit=self.unit, level=level)
 
 
 class DynamicArrayVariable(ArrayVariable):

brian2/groups/group.py

@@ -55,9 +55,11 @@ def __getitem__(self, index):
             check_code_units(abstract_code, self.group,
                              additional_variables=self.variables,
                              additional_namespace=additional_namespace)
+            template = getattr(self.group, '_index_with_code_template',
+                              'state_variable_indexing')
             codeobj = create_runner_codeobj(self.group,
                                             abstract_code,
-                                            'state_variable_indexing',
+                                            template,
                                             additional_variables=self.variables,
                                             additional_namespace=additional_namespace,
                                             )
@@ -175,7 +177,7 @@ def __setattr__(self, name, val):
             object.__setattr__(self, name, val)
 
     def _set_with_code(self, variable, group_indices, code,
-                       check_units=True, level=0):
+                       template, check_units=True, level=0):
         '''
         Sets a variable using a string expression. Is called by
         `VariableView.__setitem__` for statements such as
@@ -190,6 +192,8 @@ def _set_with_code(self, variable, group_indices, code,
         code : str
             The code that should be executed to set the variable values.
             Can contain references to indices, such as `i` or `j`
+        template : str
+            The name of the template to use.
         check_units : bool, optional
             Whether to check the units of the expression.
         level : int, optional
@@ -209,7 +213,7 @@ def _set_with_code(self, variable, group_indices, code,
         # array for situations where iterate_all could be used
         codeobj = create_runner_codeobj(self,
                                  abstract_code,
-                                 'group_variable_set',
+                                 template,
                                  additional_variables=additional_variables,
                                  additional_namespace=additional_namespace,
                                  check_units=check_units)
@@ -310,7 +314,11 @@ def create_runner_codeobj(group, code, template_name, indices=None,
         A dictionary of additional information that is passed to the template.
     '''
     logger.debug('Creating code object for abstract code:\n' + str(code))
-
+
+    if check_units:
+        check_code_units(code, group, additional_variables=additional_variables,
+                         additional_namespace=additional_namespace)
+
     template = get_codeobject_template(template_name,
                                        codeobj_class=group.codeobj_class)
 
@@ -323,15 +331,16 @@ def create_runner_codeobj(group, code, template_name, indices=None,
                                                  recursive=True)
 
     logger.debug('Unknown identifiers in the abstract code: ' + str(unknown))
-    resolved_namespace = group.namespace.resolve_all(unknown,
-                                                     additional_namespace)
 
     # Only pass the variables that are actually used
     variables = {}
     for var in used_known:
         if not isinstance(all_variables[var], StochasticVariable):
             variables[var] = all_variables[var]
 
+    resolved_namespace = group.namespace.resolve_all(unknown,
+                                                     additional_namespace)
+
     # Also add the variables that the template needs
     for var in template.variables:
         try: