Work in progress

snudda/place/place.py

@@ -570,28 +570,38 @@ def avoid_edges_parallel(self):
         ss = np.random.SeedSequence(self.random_seed + 100)
         neuron_random_seed = ss.generate_state(len(self.neurons))
 
+        bend_neuron_info = []
+
+        for neuron in self.neurons:
+            config = self.config["Neurons"][neuron.name]
+
+            if "stayInsideMesh" in config and config["stayInsideMesh"]:
+                volume_id = config["volumeID"]
+                mesh_file = self.config["Volume"][volume_id]["meshFile"]
+                bend_neuron_info.append((neuron.neuron_id, neuron.name, neuron.swc_filename,
+                                         neuron.position, neuron.rotation,
+                                         neuron_random_seed[neuron.neuron_id],
+                                         volume_id, mesh_file))
+
         if self.d_view is None:
             # Make sure we use the same random seeds if we run in serial, as would have been used in parallel
-            return self.avoid_edges(neuron_random_seeds=neuron_random_seed)
-
-        # Make random permutation of neurons, to spread out the edge neurons
-        unsorted_neuron_id = self.random_generator.permutation(len(self.neurons))
-
-        self.d_view.scatter("worker_neuron_id", unsorted_neuron_id, block=True)
-        # TODO: We cant pickle self.neurons it seems
-        import pdb
-        pdb.set_trace()
-        self.d_view.push({"config": self.config,
-                         "neurons": self.neurons,
-                         "neuron_random_seeds": neuron_random_seed}, block=True)
-
-        cmd_str = f"sp = SnuddaPlace(config_file={self.config_file},network_path={self.network_path},snudda_data={self.snudda_data}, random_seed=worker_random_seed[0])"
-        self.d_view.execute(cmd_str)
-        cmd_str2 = f"sp.config = config; sp.neurons = neurons"
-        self.d_view.execute(cmd_str2)
-        cmd_str3 = f"modified_neurons = sp.avoid_edges(neuron_id=unsorted_neuron_id, random_seed=neuron_random_seeds[unsorted_neuron_id])"
-        self.d_view.execute(cmd_str3)
-        modified_neurons = self.d_view.gather("modified_neurons", block=True)
+
+            modified_neurons = SnuddaPlace.avoid_edges_helper(bend_neuron_info=bend_neuron_info, network_path=self.network_path)
+
+        else:
+
+            # Make random permutation of neurons, to spread out the edge neurons
+            unsorted_neuron_id = self.random_generator.permutation(len(bend_neuron_info))
+            bend_neuron_info = [bend_neuron_info[idx] for idx in unsorted_neuron_id]
+
+            self.d_view.scatter("bend_neuron_info", bend_neuron_info, block=True)
+
+            cmd_str = f"sp = SnuddaPlace(config_file={self.config_file},network_path={self.network_path},snudda_data={self.snudda_data})"
+            self.d_view.execute(cmd_str)
+            cmd_str3 = f"modified_neurons = sp.avoid_edges_helper(bend_neuron_info=bend_neuron_info)"
+            self.d_view.execute(cmd_str3)
+
+            modified_neurons = self.d_view.gather("modified_neurons", block=True)
 
         for neuron_id, new_morphology in modified_neurons:
             # Replace the original morphology with the warped morphology, morphology includes rotation
@@ -600,55 +610,41 @@ def avoid_edges_parallel(self):
 
     def avoid_edges(self, neuron_id=None, neuron_random_seeds=None):
 
+        neuron_name = [n.name for n in self.neurons]
+
+        SnuddaPlace.avoid_edges_helper(neuron_id, )
+
+    def avoid_edges_helper(bend_neuron_info, network_path):
+
         # TODO: We need name, swc_file, position, rotation
         # This needs to be passed, since self.neurons is not pickleable...
 
         from snudda.place.bend_morphologies import BendMorphologies
 
         bend_morph = dict()
-        bend_morph_path = os.path.join(self.network_path, "modified_morphologies")
+        bend_morph_path = os.path.join(network_path, "modified_morphologies")
 
         if not os.path.isdir(bend_morph_path):
             os.mkdir(bend_morph_path)
 
-        if neuron_id is None:
-            neurons = self.neurons
-        else:
-            neurons = self.neurons[neuron_id]
-
-        if neuron_random_seeds is None:
-            neuron_random_seeds = [None for n in neurons]
-        else:
-            neuron_random_seeds = neuron_random_seeds[neuron_id]
-
         modified_morphologies = []
 
-        for neuron, random_seed in zip(neurons, neuron_random_seeds.flatten()):
-
-            # print(f"neuron.name = {neuron.name}, random_seed = {random_seed}")
-            config = self.config["Neurons"][neuron.name]
-
-            if "stayInsideMesh" in config and config["stayInsideMesh"]:
-                volume_id = config["volumeID"]
-
-                if volume_id not in bend_morph:
-                    mesh_file = self.config["Volume"][volume_id]["meshFile"]
-                    bend_morph[volume_id] = BendMorphologies(region_mesh=mesh_file, rng=self.random_generator)
+        for neuron_id, neuron_name, swc_filename, position, rotation, random_seed, volume_id, mesh_file \
+                in bend_neuron_info:
 
-                # Returns None if unchanged
-                new_morph_name = os.path.join(bend_morph_path, f"{neuron.name}-{neuron.neuron_id}.swc")
-                new_morphology = bend_morph[volume_id].edge_avoiding_morphology(swc_file=neuron.swc_filename,
-                                                                                new_file=new_morph_name,
-                                                                                original_position=neuron.position,
-                                                                                original_rotation=neuron.rotation,
-                                                                                random_seed=random_seed)
+            if volume_id not in bend_morph:
+                bend_morph[volume_id] = BendMorphologies(region_mesh=mesh_file, rng=None)
 
-                if new_morphology:
-                    # Replace the original morphology with the warped morphology, morphology includes rotation
-                    neuron.swc_filename = new_morphology
-                    neuron.rotation = np.eye(3)
+            # Returns None if unchanged
+            new_morph_name = os.path.join(bend_morph_path, f"{neuron_name}-{neuron_id}.swc")
+            new_morphology = bend_morph[volume_id].edge_avoiding_morphology(swc_file=swc_filename,
+                                                                            new_file=new_morph_name,
+                                                                            original_position=position,
+                                                                            original_rotation=rotation,
+                                                                            random_seed=random_seed)
 
-                    modified_morphologies.append((neuron.neuron_id, new_morphology))
+            if new_morphology:
+                modified_morphologies.append((neuron_id, new_morphology))
 
         return modified_morphologies