Merge pull request #1556 from kjohnsen/master

Fix VariableView.shape for subgroups

brian2/core/variables.py

@@ -1574,7 +1574,15 @@ def __hash__(self):
     @property
     def shape(self):
         if self.ndim == 1:
-            return (self.variable.size,)
+            if not self.variable.scalar:
+                # This is safer than using the variable size, since it also works for subgroups
+                # see GitHub issue #1555
+                size = self.group.stop - self.group.start
+                assert size <= self.variable.size
+            else:
+                size = self.variable.size
+
+            return (size,)
         else:
             return self.variable.size
 

brian2/tests/test_neurongroup.py

@@ -19,7 +19,11 @@
 from brian2.synapses.synapses import Synapses
 from brian2.tests.utils import assert_allclose, exc_isinstance
 from brian2.units.allunits import second, volt
-from brian2.units.fundamentalunits import DimensionMismatchError, have_same_dimensions
+from brian2.units.fundamentalunits import (
+    DIMENSIONLESS,
+    DimensionMismatchError,
+    have_same_dimensions,
+)
 from brian2.units.stdunits import Hz, ms, mV
 from brian2.units.unitsafefunctions import linspace
 from brian2.utils.logger import catch_logs
@@ -147,6 +151,31 @@ def test_variableview_calculations():
         2**G.y  # raising to a power with units
 
 
+@pytest.mark.standalone_compatible
+def test_variableview_properties():
+    G = NeuronGroup(
+        10,
+        """
+    x : 1
+    y : volt
+    idx : integer
+    """,
+    )
+    # The below properties should not require access to the values
+    G.x = "rand()"
+    G.y = "rand()*mV"
+    G.idx = "int(rand()*10)"
+
+    assert have_same_dimensions(G.x.unit, DIMENSIONLESS)
+    assert have_same_dimensions(G.y.unit, volt)
+    assert have_same_dimensions(G.idx.unit, DIMENSIONLESS)
+
+    assert G.x.shape == G.y.shape == G.idx.shape == (10,)
+    assert G.x.ndim == G.y.ndim == G.idx.ndim == 1
+    assert G.x.dtype == G.y.dtype == prefs.core.default_float_dtype
+    assert G.idx.dtype == np.int32
+
+
 @pytest.mark.codegen_independent
 def test_variableview_inplace_calculations():
     # Check that you can directly do in-place calculation with "variable views"

brian2/tests/test_subgroup.py

@@ -5,6 +5,7 @@
 from brian2.core.network import schedule_propagation_offset
 from brian2.devices.device import reinit_and_delete
 from brian2.tests.utils import assert_allclose
+from brian2.units.fundamentalunits import DIMENSIONLESS
 from brian2.utils.logger import catch_logs
 
 
@@ -132,6 +133,90 @@ def test_state_variables_group_as_index_problematic():
             )
 
 
+@pytest.mark.standalone_compatible
+def test_variableview_calculations():
+    # Check that you can directly calculate with "variable views"
+    G = NeuronGroup(
+        10,
+        """
+        x : 1
+        y : volt
+        idx : integer
+        """,
+    )
+    G.x = np.arange(10)
+    G.y = np.arange(10)[::-1] * mV
+    G.idx = np.arange(10, dtype=int)
+    SG = G[3:8]
+
+    assert_allclose(SG.x * SG.y, np.arange(3, 8) * np.arange(6, 1, -1) * mV)
+    assert_allclose(-SG.x, -np.arange(3, 8))
+    assert_allclose(-SG.y, -np.arange(6, 1, -1) * mV)
+
+    assert_allclose(3 * SG.x, 3 * np.arange(3, 8))
+    assert_allclose(3 * SG.y, 3 * np.arange(6, 1, -1) * mV)
+    assert_allclose(SG.x * 3, 3 * np.arange(3, 8))
+    assert_allclose(SG.y * 3, 3 * np.arange(6, 1, -1) * mV)
+    assert_allclose(SG.x / 2.0, np.arange(3, 8) / 2.0)
+    assert_allclose(SG.y / 2, np.arange(6, 1, -1) * mV / 2)
+    assert_equal(SG.idx % 2, np.arange(3, 8, dtype=int) % 2)
+    assert_allclose(SG.x + 2, 2 + np.arange(3, 8))
+    assert_allclose(SG.y + 2 * mV, 2 * mV + np.arange(6, 1, -1) * mV)
+    assert_allclose(2 + SG.x, 2 + np.arange(3, 8))
+    assert_allclose(2 * mV + SG.y, 2 * mV + np.arange(6, 1, -1) * mV)
+    assert_allclose(SG.x - 2, np.arange(3, 8) - 2)
+    assert_allclose(SG.y - 2 * mV, np.arange(6, 1, -1) * mV - 2 * mV)
+    assert_allclose(2 - SG.x, 2 - np.arange(3, 8))
+    assert_allclose(2 * mV - SG.y, 2 * mV - np.arange(6, 1, -1) * mV)
+    assert_allclose(SG.x**2, np.arange(3, 8) ** 2)
+    assert_allclose(SG.y**2, (np.arange(6, 1, -1) * mV) ** 2)
+    assert_allclose(2**SG.x, 2 ** np.arange(3, 8))
+
+    # incorrect units
+    with pytest.raises(DimensionMismatchError):
+        SG.x + SG.y
+    with pytest.raises(DimensionMismatchError):
+        SG.x[:] + SG.y
+    with pytest.raises(DimensionMismatchError):
+        SG.x + SG.y[:]
+    with pytest.raises(DimensionMismatchError):
+        SG.x + 3 * mV
+    with pytest.raises(DimensionMismatchError):
+        3 * mV + SG.x
+    with pytest.raises(DimensionMismatchError):
+        SG.y + 3
+    with pytest.raises(DimensionMismatchError):
+        3 + SG.y
+    with pytest.raises(TypeError):
+        2**SG.y  # raising to a power with units
+
+
+@pytest.mark.standalone_compatible
+def test_variableview_properties():
+    G = NeuronGroup(
+        10,
+        """
+    x : 1
+    y : volt
+    idx : integer
+    """,
+    )
+    # The below properties should not require access to the values
+    G.x = "rand()"
+    G.y = "rand()*mV"
+    G.idx = "int(rand()*10)"
+    SG = G[3:8]
+
+    assert have_same_dimensions(SG.x.unit, DIMENSIONLESS)
+    assert have_same_dimensions(SG.y.unit, volt)
+    assert have_same_dimensions(SG.idx.unit, DIMENSIONLESS)
+    # See github issue #1555
+    assert SG.x.shape == SG.y.shape == SG.idx.shape == (5,)
+    assert SG.x.ndim == SG.y.ndim == SG.idx.ndim == 1
+    assert SG.x.dtype == SG.y.dtype == prefs.core.default_float_dtype
+    assert SG.idx.dtype == np.int32
+
+
 @pytest.mark.standalone_compatible
 def test_state_monitor():
     G = NeuronGroup(10, "v : volt")