Integrating code in workflow, added stayInsideMesh option to neurons …

…(it might get renamed)

snudda/init/init.py

@@ -409,7 +409,8 @@ def add_neurons(self, name,
                     axon_config=None,
                     model_type="neuron",
                     volume_id=None,
-                    rotation_mode="random"):
+                    rotation_mode="random",
+                    stay_inside=False):
 
         if num_neurons <= 0:
             return
@@ -563,6 +564,9 @@ def add_neurons(self, name,
             if axon_config is not None:
                 cell_data["axonConfig"] = axon_config
 
+            if stay_inside:
+                cell_data["stayInsideMesh"] = True
+
             self.network_data["Neurons"][unique_name] = cell_data
 
     def get_morphologies(self, neuron_dir):

snudda/place/bend_morphologies.py

@@ -95,7 +95,7 @@ def bend_morphology_OLD(self, morphology: NeuronMorphologyExtended, k=50e-6):
 
             parent_rotation_matrices[point_idx] = rotation_matrix
 
-    def bend_morphology(self, morphology: NeuronMorphologyExtended, k=10e-6, n_random=5):
+    def bend_morphology(self, morphology: NeuronMorphologyExtended, k=30e-6, n_random=20):
 
         # k -- how early will the neuron start bending when it approaches the border
 
@@ -105,6 +105,7 @@ def bend_morphology(self, morphology: NeuronMorphologyExtended, k=10e-6, n_rando
 
         old_rotation_representation = self.get_full_rotation_representation(morphology=morphology)
         new_rotation_representation = dict()
+        morphology_changed = False
 
         for section in morphology.section_iterator():
             if (section.section_id, section.section_type) in parent_direction:
@@ -141,24 +142,30 @@ def bend_morphology(self, morphology: NeuronMorphologyExtended, k=10e-6, n_rando
 
                 if self.rng.uniform() < P_move:
 
+                    morphology_changed = True
+
                     # We need to randomize new rotation matrix
                     # https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.html
                     angles = self.rng.uniform(size=(n_random, 3), low=-0.1, high=0.1)  # Angles in radians
-                    avoidance_rotation = Rotation.from_euler(seq="XYZ", angles=angles)
+                    avoidance_rotations = Rotation.from_euler(seq="XYZ", angles=angles)
 
-                    for idx, av_rot in enumerate(avoidance_rotation):
+                    for idx, av_rot in enumerate(avoidance_rotations):
                         candidate_pos[idx, :] = parent_point + length * (av_rot*rotation).apply(vectors=parent_dir)
 
                     candidate_dist = self.region_mesh.distance_to_border(points=candidate_pos)
 
+                    #import pdb
+                    #pdb.set_trace()
+
                     if False:
                         P_candidate = np.divide(1, 1 + np.exp(-candidate_dist/k))
                         P_candidate = P_candidate / np.sum(P_candidate)
                         picked_idx = self.rng.choice(n_random, p=P_candidate)
                     else:
+                        # We want the smallest (or most negative) distance
                         picked_idx = np.argsort(candidate_dist)[0]
 
-                    new_rot = avoidance_rotation[picked_idx] * rotation
+                    new_rot = avoidance_rotations[picked_idx] * rotation
                     new_rot_rep.append((new_rot, length))
                     segment_direction = new_rot.apply(parent_dir)
 
@@ -175,7 +182,7 @@ def bend_morphology(self, morphology: NeuronMorphologyExtended, k=10e-6, n_rando
 
             new_rotation_representation[section.section_id, section.section_type] = new_rot_rep
 
-        return new_rotation_representation
+        return new_rotation_representation, morphology_changed
 
     def get_full_rotation_representation(self, morphology: MorphologyData):
 
@@ -351,6 +358,19 @@ def write_swc(self, morphology: MorphologyData, output_file, comment=None):
 
         print(f"Wrote {output_file}")
 
+    def edge_avoiding_morphology(self, swc_file, new_file):
+
+        md = MorphologyData(swc_file=swc_file)
+        rot_rep, morphology_changed = self.bend_morphology(md)
+
+        if morphology_changed:
+            new_coord = self.apply_rotation(md, rot_rep)
+            md.geometry[:, :3] = new_coord
+            self.write_swc(morphology=md, swc_file=new_file)
+            return new_file
+
+        # Returns None if morphology was not changed
+        return None
 
 def test_rotation_representation():
 
@@ -399,7 +419,7 @@ def test_bending():
     before = nm.clone(position=pos, rotation=np.eye(3))
     after = nm.clone(position=pos, rotation=np.eye(3))
 
-    new_rot_rep = bm.bend_morphology(after.get_morphology())
+    new_rot_rep, _ = bm.bend_morphology(after.get_morphology())
     new_coord = bm.apply_rotation(after.get_morphology(), new_rot_rep)
     after.get_morphology().geometry[:, :3] = new_coord
 

snudda/place/place.py

@@ -91,6 +91,8 @@ def __init__(self,
 
         self.network_path = network_path
         self.config_file = config_file
+        self.config = None
+
         self.axon_config_cache = None
 
         self.snudda_data = get_snudda_data(snudda_data=snudda_data,
@@ -152,6 +154,7 @@ def place(self):
         """ Place neurons in 3D space. """
 
         self.parse_config()
+        self.avoid_edges()
         self.write_data()
 
     ############################################################################
@@ -544,6 +547,7 @@ def parse_config(self, config_file=None, resort_neurons=True):
                              config=config)
 
         self.config_file = config_file
+        self.config = config
 
         # We reorder neurons, sorting their IDs after position
         # -- UPDATE: Now we spatial cluster neurons depending on number of workers
@@ -560,6 +564,38 @@ def parse_config(self, config_file=None, resort_neurons=True):
 
     ############################################################################
 
+    def avoid_edges(self):
+
+        from snudda.place.bend_morphologies import BendMorphologies
+
+        bend_morph = dict()
+        bend_morph_path = os.path.join(self.network_path, "modified_morphologies")
+
+        if not os.path.isdir(bend_morph_path):
+            os.mkdir(bend_morph_path)
+
+        for neuron in self.neurons:
+            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)
+
+                # Returns None if unchanged
+                new_morph_name = os.path.join(bend_morph_path, f"{neuron.name}-{neuron.neuron_id}")
+                new_morphology = bend_morph[volume_id].edge_avoiding_morphology(swc_file=neuron.swc_filename,
+                                                                                new_file=new_morph_name)
+
+                if new_morphology:
+                    # Replace the original morphology with the warped morphology, morphology includes rotation
+                    neuron.swc_filename = new_morphology
+                    neuron.rotation = np.eye(3)
+
+    ############################################################################
+
     def generate_extra_axon_info(self, source_neuron, position, config, rng):
 
         axon_info = []