Merge branch 'main' of https://github.com/GazzolaLab/MiV-Simulator

Author: iraikov

.pre-commit-config.yaml

@@ -1,12 +1,12 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: trailing-whitespace
       - id: check-json
       - id: end-of-file-fixer
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.7
+    rev: v0.12.9
     hooks:
       - id: ruff
         args: [--fix]

src/miv_simulator/env.py

@@ -58,7 +58,7 @@ def __init__(
         config: Optional[str] = None,
         template_paths: str = "templates",
         hoc_lib_path: Optional[str] = None,
-        mechanisms_path: Optional[str] = None,
+        mechanisms_path: Optional[str] = "mechanisms",
         dataset_prefix: Optional[str] = None,
         results_path: Optional[str] = None,
         results_file_id: Optional[str] = None,

src/miv_simulator/interface/synapses.py

@@ -70,6 +70,10 @@ def __call__(self):
     def on_write_meta_data(self):
         return MPI.COMM_WORLD.Get_rank() == 0
 
+    def on_after_dispatch(self, success):
+        # ensure shutdown
+        MPI.COMM_WORLD.Abort(int(not success))
+
     def compute_context(self):
         context = super().compute_context()
         del context["config"]["forest_filepath"]

src/miv_simulator/optimization.py

@@ -377,27 +377,21 @@ def network_features(env, t_start, t_stop, target_populations):
 
     for pop_name in target_populations:
         n_active = 0
-        sum_mean_rate = 0.0
         spike_density_dict = spikedata.spike_density_estimate(
             pop_name, pop_spike_dict[pop_name], time_bins
         )
         for gid, dens_dict in spike_density_dict.items():
             mean_rate = np.mean(dens_dict["rate"])
-            sum_mean_rate += mean_rate
             if mean_rate > 0.0:
                 n_active += 1
 
         n_total = len(env.cells[pop_name]) - len(env.artificial_cells[pop_name])
 
-        n_target_rate_map = 0
-        sum_snr = None
-
         pop_features_dict = {}
         pop_features_dict["n_total"] = n_total
         pop_features_dict["n_active"] = n_active
-        pop_features_dict["n_target_rate_map"] = n_target_rate_map
-        pop_features_dict["sum_mean_rate"] = sum_mean_rate
-        pop_features_dict["sum_snr"] = sum_snr
+        pop_features_dict["time_bins"] = time_bins
+        pop_features_dict["spike_density_dict"] = spike_density_dict
 
         features_dict[pop_name] = pop_features_dict
 

src/miv_simulator/optimize_network.py

@@ -374,57 +374,80 @@ def compute_objectives(local_features, operational_config, opt_targets):
     all_features_dict = {}
     constraints = []
 
+    active_threshold = 0.01
     target_populations = operational_config["target_populations"]
+    temporal_resolution = operational_config["temporal_resolution"]
     for pop_name in target_populations:
         pop_features_dicts = [
             features_dict[0][pop_name] for features_dict in local_features
         ]
 
         sum_mean_rate = 0.0
-        sum_snr = 0.0
         n_total = 0
         n_active = 0
-        n_target_rate_map = 0
+        time_bins_ref = None
+        sum_active_per_bin = None
         for pop_feature_dict in pop_features_dicts:
             n_active_local = pop_feature_dict["n_active"]
             n_total_local = pop_feature_dict["n_total"]
-            n_target_rate_map_local = pop_feature_dict["n_target_rate_map"]
-            sum_mean_rate_local = pop_feature_dict["sum_mean_rate"]
-            sum_snr_local = pop_feature_dict["sum_snr"]
+            time_bins = pop_feature_dict["time_bins"]
+            if time_bins_ref is None:
+                time_bins_ref = time_bins
+            spike_density_dict = pop_feature_dict["spike_density_dict"]
+            sum_mean_rate_local = 0.0
+            t_start = time_bins_ref[0]
+            t_end = time_bins_ref[-1] + (time_bins_ref[1] - time_bins_ref[0])
+            # time bins for fraction active per time bin calculation
+            fr_time_bins = np.arange(t_start, t_end, temporal_resolution)
+            bin_width = time_bins_ref[1] - time_bins_ref[0]
+            time_centers = time_bins_ref + bin_width / 2
+            fr_time_centers = fr_time_bins + temporal_resolution / 2
+            if sum_active_per_bin is None:
+                sum_active_per_bin = np.zeros_like(time_centers)
+            for gid, dens_dict in spike_density_dict.items():
+                mean_rate = np.mean(dens_dict["rate"])
+                if mean_rate > 0.0:
+                    sum_mean_rate_local += mean_rate
+                ip_rate = np.interp1d(
+                    fr_time_centers,
+                    dens_dict["rate"],
+                    kind="linear",
+                    bounds_error=False,
+                    fill_value=0.0,
+                )
+                active_per_bin = ip_rate > active_threshold
+                sum_active_per_bin += active_per_bin
 
             n_total += n_total_local
             n_active += n_active_local
-            n_target_rate_map += n_target_rate_map_local
             sum_mean_rate += sum_mean_rate_local
 
-            if sum_snr_local is not None:
-                sum_snr += sum_snr_local
-
         if n_active > 0:
             mean_rate = sum_mean_rate / n_active
         else:
             mean_rate = 0.0
 
         if n_total > 0:
             fraction_active = n_active / n_total
+            mean_fraction_active_per_bin = np.mean(sum_active_per_bin / float(n_total))
+            std_fraction_active_per_bin = np.std(sum_active_per_bin / float(n_total))
         else:
             fraction_active = 0.0
-
-        mean_snr = None
-        if n_target_rate_map > 0:
-            mean_snr = sum_snr / n_target_rate_map
+            mean_fraction_active_per_bin = 0.0
+            std_fraction_active_per_bin = 0.0
 
         logger.info(
             f"population {pop_name}: n_active = {n_active} n_total = {n_total} mean rate = {mean_rate}"
         )
-        logger.info(
-            f"population {pop_name}: n_target_rate_map = {n_target_rate_map} snr: sum = {sum_snr} mean = {mean_snr}"
-        )
 
+        all_features_dict[f"{pop_name} mean fraction active per time bin"] = (
+            mean_fraction_active_per_bin
+        )
+        all_features_dict[f"{pop_name} std fraction active per time bin"] = (
+            std_fraction_active_per_bin
+        )
         all_features_dict[f"{pop_name} fraction active"] = fraction_active
         all_features_dict[f"{pop_name} firing rate"] = mean_rate
-        if mean_snr is not None:
-            all_features_dict[f"{pop_name} snr"] = mean_snr
 
         rate_constr = mean_rate if mean_rate > 0.0 else -1.0
         constraints.append(rate_constr)

src/miv_simulator/scripts/distribute_synapse_locs.py

@@ -2,6 +2,7 @@
 
 
 import click
+from mpi4py import MPI
 from miv_simulator.simulator.distribute_synapses import (
     distribute_synapse_locations,
 )
@@ -75,3 +76,5 @@ def main(
         config_prefix=config_prefix,
         mechanisms_path=mechanisms_path,
     )
+
+    MPI.COMM_WORLD.Abort(0)

src/miv_simulator/simulator/distribute_synapses.py

@@ -264,7 +264,7 @@ def distribute_synapses(
     dry_run: bool,
 ):
     logger = utils.get_script_logger(os.path.basename(__file__))
-
+    write_size = 1
     comm = MPI.COMM_WORLD
     rank = comm.rank
 
@@ -413,7 +413,15 @@ def distribute_synapses(
             else:
                 logger.info(f"Rank {rank} gid is None")
             gc.collect()
-            if (not dry_run) and (write_size > 0) and (gid_count % write_size == 0):
+
+            # Check if any rank has reached the write_size threshold and
+            #  ensure that all ranks in that group call collectively
+            local_should_write = (
+                (not dry_run) and (write_size > 0) and (gid_count % write_size == 0)
+            )
+            global_should_write = comm.allreduce(local_should_write, op=MPI.LOR)
+
+            if global_should_write:
                 append_cell_attributes(
                     output_filepath,
                     population,
@@ -428,7 +436,13 @@ def distribute_synapses(
             syn_stats[population] = syn_stats_dict
             count += 1
 
-        if not dry_run:
+        # Final write for any remaining synapse data - allreduce to ensure all ranks participate
+        local_should_write_final = not dry_run
+        global_should_write_final = comm.allreduce(
+            local_should_write_final, op=MPI.LAND
+        )
+
+        if global_should_write_final:
             append_cell_attributes(
                 output_filepath,
                 population,
@@ -447,7 +461,3 @@ def distribute_synapses(
                 f"to compute synapse locations for {np.sum(global_count)} cells"
             )
             logger.info(summary)
-
-        comm.barrier()
-
-    MPI.Finalize()

src/miv_simulator/synapses.py

@@ -2425,10 +2425,11 @@ def modify_mechanism_parameters(
 
             assert new_val is not None
 
-            value_result, source = (
-                self.syn_store.param_store.get_parameter_value_hierarchy(
-                    gid, array_index, mech_name, k, syn_id
-                )
+            (
+                value_result,
+                source,
+            ) = self.syn_store.param_store.get_parameter_value_hierarchy(
+                gid, array_index, mech_name, k, syn_id
             )
             old_val = value_result if value_result is not None else mech_param
 
@@ -3133,13 +3134,12 @@ def insert_cell_syns(
                             list(sec_list),
                         )
                     )
-        if len(reduced_section_dict) == 0:
-            if hasattr(cell, "soma"):
-                cell_soma = cell.soma
-                if isinstance(cell_soma, list):
-                    cell_soma = cell_soma[0]
-            if hasattr(cell, "dend"):
-                cell_dendrite = cell.dend
+        if hasattr(cell, "soma"):
+            cell_soma = cell.soma
+            if isinstance(cell_soma, list):
+                cell_soma = cell_soma[0]
+        if hasattr(cell, "dend"):
+            cell_dendrite = cell.dend
 
     syn_manager = env.synapse_manager