Merged in loaddistributedexp (pull request #8)

[NRPJP-83] errors in loading experiments were due to incorrectly parsing populations attached to devices. This commit fixes it.

* [NRPJP-83] errors in loading experiments were due to incorrectly parsing populations attached to devices. This commit fixes it.

* [NRRPLT-83] commented out annoying print

* [NRRPLT-83] Added suggestions from reviewer and removed unnecessary commented code

* [NRRPLT-83] Rearranged some code into a function + added function descriptions

* [NRRPLT-83] Fixed bugs introduced along PR

* [NRRPLT-83] Fixed index_from_assembly not considering that populations can be instances of Population

Approved-by: Kepa Cantero <cantero@fortiss.org>
Approved-by: Ugo Albanese <ugo.albanese@santannapisa.it>

hbp_nrp_distributed_nest/hbp_nrp_distributed_nest/cle/DistributedPyNNCommunicationAdapter.py

@@ -33,6 +33,8 @@ import hbp_nrp_cle.tf_framework.config as tf_config
 
 from mpi4py import MPI
 
+import pyNN.nest as sim
+
 import logging
 import time
 
@@ -70,8 +72,12 @@ class DistributedPyNNCommunicationAdapter(PyNNNestCommunicationAdapter):
         device.timestep = ts
 
         # mark the device as MPI-aware, only used by Nest-specific devices
-        if isinstance(device, PyNNNestDevice):
+        if not isinstance(populations, list) and isinstance(device, PyNNNestDevice):
             setattr(device, 'mpi_aware', True)
+        else:
+            for d in device.devices:
+                if isinstance(d, PyNNNestDevice):
+                    setattr(d, 'mpi_aware', True)
         return device
 
     def register_spike_sink(self, populations, spike_detector_type, **params):
@@ -135,18 +141,42 @@ class DistributedPyNNCommunicationAdapter(PyNNNestCommunicationAdapter):
 
         # population information to send to the remote MPI nodes, we can't pickle Populations
         # directly and those references wouldn't be valid on the remote nodes anyway
-        if populations.label in tf_config.brain_root.__dict__:
-            label = populations.label
-        else:
-            label = populations.parent.label
-        mask = populations.mask if populations.mask else slice(0, len(populations), 1)
+
+        # create a structure describing 'populations'
+        br = tf_config.brain_root
+        assemblies = []
+        if not isinstance(populations, list):
+            populations = [populations]
+
+        def index_from_population_view(p):
+            """
+            Returns neuron indices of a PopulationView in the root Population and the root label
+            """
+            label = p.grandparent.label
+            indices = p.index_in_grandparent(range(p.size))
+            return label, indices
+
+        def index_from_assembly(a):
+            """
+            Returns neuron indices and labels for each population in an Assembly
+            """
+            return [index_from_population_view(p) if isinstance(p, sim.PopulationView) else
+                    (p.label, None) for p in a.populations]
+
+        for population in populations:
+            if isinstance(population, sim.Population):
+                assemblies += [[(population.label, None)]]
+            elif isinstance(population, sim.PopulationView):
+                assemblies += [[index_from_population_view(population)]]
+            elif isinstance(population, sim.Assembly):
+                assemblies += [index_from_assembly(population)]
 
         # propagate the synapse creation parameters to all remote notes, they will run the same
         # connection/creation commands after receiving these messages, guaranteed to be
         # run from CLE MPI process 0 only
         for rank in xrange(1, MPI.COMM_WORLD.Get_size()):
-            MPI.COMM_WORLD.send({'command': 'ConnectTF', 'type': kind, 'label': label,
-                                 'mask': mask, 'device': device, 'timestep': timestep,
+            MPI.COMM_WORLD.send({'command': 'ConnectTF', 'type': kind, 'assemblies': assemblies,
+                                 'device': device, 'timestep': timestep,
                                  'params': params},
                                 dest=rank, tag=NestBrainProcess.MPI_MSG_TAG)
 

hbp_nrp_distributed_nest/hbp_nrp_distributed_nest/launch/NestBrainProcess.py

@@ -137,7 +137,7 @@ class NestBrainProcess(object):
             elif data == 'step':
 
                 # run the coordinated simulation step
-                print "[MPI] ===================== step ======================="
+                #print "[MPI] ===================== step ======================="
                 self._brain_controller.run_step(self._timestep * 1000.0)  # msec
                 self._brain_communicator.refresh_buffers(0.0)
 
@@ -173,11 +173,54 @@ class NestBrainProcess(object):
         :param params The connectivity/synapse parameters passed by the CLE.
         """
 
-        # get the population of neurons from our dictionary, we can guarantee this is a valid
-        # population that has been declared in the BIBI at this point as the CLE will validate
-        # before sending us the create message
-        brain_pop = sim.PopulationView(tf_config.brain_root.__dict__[params['label']],
-                                       selector=params['mask'])
+        # get the population of neurons from params
+
+        def get_population_from_label(pop, label):
+            """
+            Given an input population and a label returns a population with the specified label.
+            This can be the population itself or its root (grandparent)
+            """
+            if pop.label == label:
+                return pop
+            elif isinstance(pop, sim.PopulationView) and pop.grandparent.label == label:
+                return pop.grandparent
+            else:
+                return None
+
+        def get_population_from_brain(label):
+            """
+            Finds a population in brain with an specified label.
+            """
+
+            # check if label in brain
+            if label in tf_config.brain_root.__dict__:
+                return tf_config.brain_root.__dict__[label]
+            # check if label in circuit populations
+            elif 'circuit' in tf_config.brain_root.__dict__:
+                circuit = tf_config.brain_root.__dict__['circuit']
+                pops = circuit.populations if isinstance(circuit,sim.Assembly) else [circuit]
+                for pop in pops:
+                    p = get_population_from_label(pop,label)
+                    if p:
+                        return p
+            # could not find population
+            return None
+
+        populations = []
+        for assembly in params['assemblies']:
+            a = []
+            for p in assembly:
+                population = get_population_from_brain(p[0])
+                if population and p[1] is not None:
+                    population = sim.PopulationView(population,p[1])
+                a += [population]
+
+            if len(a) == 1:
+                populations += [a[0]]
+            else:
+                populations += [sim.Assembly(*a)]
+
+        brain_pop = populations[0] if len(populations) == 1 else populations
 
         # perform the actual device creation/connection
         if params['type'] == 'sink':
@@ -193,8 +236,12 @@ class NestBrainProcess(object):
         device.timestep = params['timestep']
 
         # mark the device as MPI-aware, only used by Nest-specific devices
-        if isinstance(device, PyNNNestDevice):
+        if not isinstance(brain_pop,list) and isinstance(device, PyNNNestDevice):
             setattr(device, 'mpi_aware', True)
+        else:
+            for d in device.devices:
+                if isinstance(d, PyNNNestDevice):
+                    setattr(d, 'mpi_aware', True)
 
         # store the device in a way that we can easily retrieve later
         self.devices[device.timestep] = device