Update network interface

GazzolaLab · Jun 29, 2024 · a5bb6b9 · a5bb6b9
1 parent b415459
commit a5bb6b9
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 18 deletions.
diff --git a/src/miv_simulator/connections.py b/src/miv_simulator/connections.py
@@ -649,10 +649,10 @@ def generate_uv_distance_connections(
 
             for source_population in source_populations:
                 source_layers = projection_config[source_population].layers
-                projection_prob_dict[source_population] = (
-                    connection_prob.get_prob(
-                        destination_gid, source_population, source_layers
-                    )
+                projection_prob_dict[
+                    source_population
+                ] = connection_prob.get_prob(
+                    destination_gid, source_population, source_layers
                 )
 
                 for layer, (

diff --git a/src/miv_simulator/interface/architecture.py b/src/miv_simulator/interface/architecture.py
@@ -45,7 +45,7 @@ class Config(BaseModel):
         io_size: int = -1
         chunk_size: int = 1000
         value_chunk_size: int = 1000
-        ranks: int = 8
+        ranks: int = -1
         nodes: str = "1"
 
     def config_from_file(self, filename: str) -> Dict:

diff --git a/src/miv_simulator/interface/connections.py b/src/miv_simulator/interface/connections.py
@@ -31,7 +31,7 @@ class Config(BaseModel):
         value_chunk_size: int = 1000
         cache_size: int = 1
         write_size: int = 1
-        ranks: int = 8
+        ranks: int = -1
         nodes: str = "1"
 
     def config_from_file(self, filename: str) -> Dict:
@@ -69,7 +69,7 @@ def __call__(self):
             cache_size=self.config.cache_size,
             write_size=self.config.write_size,
             dry_run=False,
-            seeds=self.seed,
+            seeds=[18000000, 2000000, 1500000],
         )
 
     def on_write_meta_data(self):

diff --git a/src/miv_simulator/interface/distances.py b/src/miv_simulator/interface/distances.py
@@ -28,7 +28,7 @@ class Config(BaseModel):
         chunk_size: int = 1000
         value_chunk_size: int = 1000
         cache_size: int = 50
-        ranks: int = 8
+        ranks: int = -1
         nodes: str = "1"
 
     def __call__(self):

diff --git a/src/miv_simulator/interface/network.py b/src/miv_simulator/interface/network.py
@@ -1,11 +1,20 @@
 import os
 
+from typing import Optional
+
 from pydantic import BaseModel, Field, ConfigDict
 from machinable import Interface, get
 from miv_simulator.config import Config
 from miv_simulator import mechanisms
 
 
+def _lp(x1, x2, y1, y2, x) -> int:
+    q = ((y2 - y1) * x + x2 * y1 - x1 * y2) / (x2 - x1)
+    q = max(x1, q)
+    q = min(x2, q)
+    return int(q)
+
+
 class Network(Interface):
     class Config(BaseModel):
         model_config = ConfigDict(extra="forbid")
@@ -14,12 +23,17 @@ class Config(BaseModel):
         mechanisms_path: str = ("./mechanisms",)
         template_path: str = ("./templates",)
         morphology_path: str = "./morphology"
+        populations: Optional[list[str]] = None
 
     def launch(self):
         self.source_config = config = Config.from_yaml(
             self.config.config_filepath
         )
 
+        populations = self.config.populations
+        if populations is None:
+            populations = list(config.synapses.keys())
+
         self.h5_types = get(
             "miv_simulator.interface.h5_types",
             [
@@ -58,7 +72,7 @@ def launch(self):
                 },
                 uses=self.distances,
             ).launch()
-            for population in config.synapses
+            for population in [p for p in config.synapses if p in populations]
         }
 
         self.synapses = {
@@ -73,12 +87,48 @@ def launch(self):
                     "distribution": "poisson",
                     "mechanisms_path": self.config.mechanisms_path,
                     "template_path": self.config.template_path,
-                    "io_size": 1,
-                    "write_size": 0,
+                    # apply heuristic based on number of cells
+                    "io_size": _lp(
+                        0,
+                        5e5,
+                        1,
+                        30,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "write_size": _lp(
+                        0,
+                        5e5,
+                        1,
+                        100,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "chunk_size": _lp(
+                        0,
+                        5e5,
+                        1000,
+                        10000,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "value_chunk_size": _lp(
+                        0,
+                        5e5,
+                        1000,
+                        200000,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "nodes": str(
+                        _lp(
+                            0,
+                            5e5,
+                            1,
+                            25,
+                            sum(config.cell_distributions[population].values()),
+                        )
+                    ),
                 },
                 uses=self.synapse_forest[population],
             ).launch()
-            for population in config.synapses
+            for population in self.synapse_forest
         }
 
         self.connections = {
@@ -91,9 +141,50 @@ def launch(self):
                     "axon_extents": config.axon_extents,
                     "population_definitions": config.definitions.populations,
                     "layer_definitions": config.definitions.layers,
-                    "io_size": 1,
-                    "cache_size": 20,
-                    "write_size": 100,
+                    "io_size": _lp(
+                        0,
+                        5e5,
+                        1,
+                        40,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "cache_size": _lp(
+                        0,
+                        5e5,
+                        1,
+                        20,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "write_size": _lp(
+                        0,
+                        5e5,
+                        1,
+                        250,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "chunk_size": _lp(
+                        0,
+                        5e5,
+                        1000,
+                        10000,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "value_chunk_size": _lp(
+                        0,
+                        5e5,
+                        1000,
+                        640000,
+                        sum(config.cell_distributions[population].values()),
+                    ),
+                    "nodes": str(
+                        _lp(
+                            0,
+                            5e5,
+                            1,
+                            64,
+                            sum(config.cell_distributions[population].values()),
+                        )
+                    ),
                 },
                 uses=self.synapses[population],
             ).launch()

diff --git a/src/miv_simulator/interface/neuroh5_graph.py b/src/miv_simulator/interface/neuroh5_graph.py
@@ -63,7 +63,7 @@ def __call__(self) -> None:
         self.graph.import_h5types(self.architecture.config.filepath)
         self.graph.import_soma_coordinates(
             self.architecture.config.filepath,
-            populations=list(populations) + ["STIM"],
+            populations=list(populations),
         )
         for p in self.synapse_forest.keys():
             self.graph.import_synapse_attributes(

diff --git a/src/miv_simulator/interface/synapses.py b/src/miv_simulator/interface/synapses.py
@@ -25,7 +25,7 @@ class Config(BaseModel):
         write_size: int = 1
         chunk_size: int = 1000
         value_chunk_size: int = 1000
-        ranks: int = 8
+        ranks: int = -1
         nodes: str = "1"
 
         @field_validator("cell_types")
@@ -63,7 +63,7 @@ def __call__(self):
             write_size=self.config.write_size,
             chunk_size=self.config.chunk_size,
             value_chunk_size=self.config.value_chunk_size,
-            seed=self.seed,
+            seed=None,
             dry_run=False,
         )