Fix population offset computation

GazzolaLab · May 31, 2024 · 8a89d17 · 8a89d17
1 parent 0d6e856
commit 8a89d17
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 16 deletions.
diff --git a/src/miv_simulator/connections.py b/src/miv_simulator/connections.py
@@ -283,8 +283,10 @@ def choose_synapse_projection(
             for ord_index in ord_indices:
                 if syn_layer == syn_config_layers[ord_index]:
                     if k not in population_dict:
-                        raise RuntimeError(f"Key {k} is not in population_dict; "
-                                           f" population_dict is {population_dict}")
+                        raise RuntimeError(
+                            f"Key {k} is not in population_dict; "
+                            f" population_dict is {population_dict}"
+                        )
                     projection_lst.append(population_dict[k])
                     projection_prob_lst.append(
                         syn_config_proportions[ord_index]
@@ -731,12 +733,11 @@ def generate_uv_distance_connections(
         projection_dict = {}
     if not dry_run:
         append_graph(
-            connectivity_path, 
-            projection_dict, 
-            io_size=io_size, 
+            connectivity_path,
+            projection_dict,
+            io_size=io_size,
             comm=comm,
             chunk_size=value_chunk_size,
-
         )
         if rank == 0:
             if connection_dict:

diff --git a/src/miv_simulator/simulator/generate_network_architecture.py b/src/miv_simulator/simulator/generate_network_architecture.py
@@ -352,8 +352,7 @@ def generate_network_architecture(
                     callback = getattr(module, obj_name)
 
                     nodes = callback(count, layer_extents[layer])
-
-                    if not len(nodes) == count:
+                    if len(nodes) != count:
                         logger.error(
                             f"Generator {layer} produced mismatch between actual count {len(nodes)} and configured count {count}"
                         )
@@ -488,6 +487,7 @@ def generate_network_architecture(
             if i % size == rank:
                 uvl_coords = all_uvl_coords_interp[coord_ind, :].ravel()
                 xyz_coords1 = all_xyz_coords_interp[coord_ind, :].ravel()
+
                 if uvl_in_bounds(
                     all_uvl_coords_interp[coord_ind, :],
                     layer_extents,
@@ -539,7 +539,6 @@ def generate_network_architecture(
 
         total_xyz_error = np.zeros((3,))
         comm.Allreduce(xyz_error, total_xyz_error, op=MPI.SUM)
-
         coords_count = 0
         coords_count = np.sum(np.asarray(comm.allgather(len(coords))))
 
@@ -555,6 +554,7 @@ def generate_network_architecture(
             )
 
         pop_coords_dict[population] = coords
+
         coords_offset += gen_coords_count
 
         if rank == 0:
@@ -644,7 +644,6 @@ def generate_network_architecture(
                                     coord_l,
                                 )
                             )
-
             sampled_coords = random_subset(all_coords, int(pop_count))
             sampled_coords.sort(
                 key=lambda coord: coord[3]

diff --git a/src/miv_simulator/simulator/generate_synapse_forest.py b/src/miv_simulator/simulator/generate_synapse_forest.py
@@ -71,12 +71,14 @@ def generate_synapse_forest(
     if not os.path.isfile(output_filepath):
         # determine population ranges
         with h5py.File(filepath, "r") as f:
-            idx = list(
-                reversed(
-                    f["H5Types"]["Population labels"].dtype.metadata["enum"]
-                )
-            ).index(population)
-            offset = f["H5Types"]["Populations"][idx][0]
+            h5type_num = f["H5Types"]["Population labels"].dtype.metadata[
+                "enum"
+            ][population]
+            population_range = [
+                p for p in f["H5Types"]["Populations"] if p[2] == h5type_num
+            ]
+            assert len(population_range) == 1
+            offset = population_range[0][0]
 
         _bin_check("neurotrees_copy")
         _run(