Perform direct conversion of allocation dict to numpy array, added ra…

…nk check

neurodamus/io/cell_readers.py

@@ -128,8 +128,6 @@ def load_nodes_base_info():
             return gidvec, meinfos, total_cells
 
         log_verbose("Loading nodes info")
-        if not isinstance(gidvec, np.ndarray):
-            gidvec = np.array(gidvec, dtype="uint32")
         node_sel = libsonata.Selection(gidvec - 1)  # 0-based node indices
         morpho_names = node_pop.get_attribute("morphology", node_sel)
         mtypes = node_pop.get_attribute("mtype", node_sel)

neurodamus/node.py

@@ -378,7 +378,17 @@ def compute_load_balance(self):
             return None
         elif lb_mode == LoadBalanceMode.Memory:
             logging.info("Load Balancing ENABLED. Mode: Memory")
-            return import_allocation_stats("allocation.pkl.gz")
+            alloc = import_allocation_stats("allocation.pkl.gz")
+            if MPI.rank == 0:
+                unique_ranks = set(rank for pop in alloc.values() for rank in pop.keys())
+                logging.debug("Unique ranks in allocation file: %s", len(unique_ranks))
+                if MPI.size != len(unique_ranks):
+                    raise ConfigurationError(
+                        "The number of ranks in the allocation file is different from the number of "
+                        "ranks in the current run. The allocation file was created with a different "
+                        "number of ranks."
+                    )
+            return alloc
 
         # Build load balancer as per requested options
         # Compat Note:

neurodamus/utils/memory.py

@@ -200,7 +200,6 @@ def export_allocation_stats(rank_allocation, filename):
         f.write(compressed_data)
 
 
-@run_only_rank0
 def import_allocation_stats(filename):
     """
     Import allocation dictionary from serialized pickle file.
@@ -209,7 +208,7 @@ def convert_to_standard_types(obj):
         """Converts an object containing defaultdicts of Vectors to standard Python types."""
         result = {}
         for node, vectors in obj.items():
-            result[node] = {key: list(vector) for key, vector in vectors.items()}
+            result[node] = {key: np.array(vector) for key, vector in vectors.items()}
         return result
 
     with open(filename, 'rb') as f: