From 6c8024502b4db65948b688d13958ee1b5a3e3f4b Mon Sep 17 00:00:00 2001
From: Ivan Raikov <ivan.g.raikov@gmail.com>
Date: Thu, 4 Apr 2024 15:52:24 -0500
Subject: [PATCH] bug fixes in check_synapses

---
 .../simulator/distribute_synapses.py          | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/miv_simulator/simulator/distribute_synapses.py b/src/miv_simulator/simulator/distribute_synapses.py
index 38074ac..b6c8679 100644
--- a/src/miv_simulator/simulator/distribute_synapses.py
+++ b/src/miv_simulator/simulator/distribute_synapses.py
@@ -15,7 +15,7 @@
     append_cell_attributes,
     read_population_ranges,
 )
-from typing import Optional, Tuple, Literal
+from typing import Optional, Tuple, Literal, Dict
 
 sys_excepthook = sys.excepthook
 
@@ -132,6 +132,7 @@ def check_syns(
         seg_density_per_sec,
         layer_set_dict,
         swc_set_dict,
+        env.layers,
         logger,
     )
 
@@ -143,30 +144,39 @@ def check_synapses(
     seg_density_per_sec,
     layer_set_dict,
     swc_set_dict,
+    swc_defs,
+    layer_defs,
     logger,
 ):
     layer_stats = syn_stats_dict["layer"]
     swc_stats = syn_stats_dict["swc_type"]
 
     warning_flag = False
+    incomplete_layers = []
     for syn_type, layer_set in list(layer_set_dict.items()):
         for layer in layer_set:
-            if layer in layer_stats:
-                if layer_stats[layer][syn_type] <= 0:
+            layer_index = layer_defs[layer]
+            if layer_index in layer_stats:
+                if layer_stats[layer_index][syn_type] <= 0:
+                    incomplete_layers.append(layer)
                     warning_flag = True
             else:
+                incomplete_layers.append(layer)
                 warning_flag = True
     if warning_flag:
         logger.warning(
-            f"Rank {MPI.COMM_WORLD.Get_rank()}: incomplete synapse layer set for cell {gid}: {layer_stats}"
+            f"Rank {MPI.COMM_WORLD.Get_rank()}: incomplete synapse layer set for cell {gid}: "
+            f"  incomplete layers: {incomplete_layers}\n"
+            f"  populated layers: {layer_stats}\n"
             f"  layer_set_dict: {layer_set_dict}\n"
             f"  seg_density_per_sec: {seg_density_per_sec}\n"
             f"  morph_dict: {morph_dict}"
         )
     for syn_type, swc_set in swc_set_dict.items():
         for swc_type in swc_set:
-            if swc_type in swc_stats:
-                if swc_stats[swc_type][syn_type] <= 0:
+            swc_type_index = swc_defs[swc_type]
+            if swc_type_index in swc_stats:
+                if swc_stats[swc_type_index][syn_type] <= 0:
                     warning_flag = True
             else:
                 warning_flag = True
@@ -222,6 +232,8 @@ def distribute_synapse_locations(
     return distribute_synapses(
         forest_filepath=forest_path,
         cell_types=env.celltypes,
+        swc_defs=env.SWC_Types,
+        layer_defs=env.layers,
         populations=populations,
         distribution=distribution,
         template_path=template_path,
@@ -238,6 +250,8 @@ def distribute_synapse_locations(
 def distribute_synapses(
     forest_filepath: str,
     cell_types: config.CellTypes,
+    swc_defs: Dict[str, int],
+    layer_defs: Dict[str, int],
     populations: Tuple[str, ...],
     distribution: Literal["uniform", "poisson"],
     template_path: str,
@@ -388,6 +402,8 @@ def distribute_synapses(
                     seg_density_per_sec,
                     layer_set_dict,
                     swc_set_dict,
+                    swc_defs,
+                    layer_defs,
                     logger,
                 )