generate_connections: use population dict from Env

GazzolaLab · Apr 5, 2024 · 1fd8aaa · 1fd8aaa
1 parent 5b0c33d
commit 1fd8aaa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 33 deletions.
diff --git a/src/miv_simulator/config.py b/src/miv_simulator/config.py
@@ -66,36 +66,11 @@ class LayersDef(IntEnum):
     SLM = 9  # Lacunosum-moleculare
 
 
-LayersLiteral = Literal[
-    "default",
-    "Hilus",
-    "GCL",
-    "IML",
-    "MML",
-    "OML",
-    "SO",
-    "SP",
-    "SL",
-    "SR",
-    "SLM",
-]
-
-
 class InputSelectivityTypesDef(IntEnum):
     random = 0
     constant = 1
 
 
-class PopulationsDef(IntEnum):
-    STIM = 0  # Stimulus
-    PYR = 100  # PYR distal dendrites
-    PVBC = 101  # Basket cells expressing parvalbumin
-    OLM = 102  # GABAergic oriens-lacunosum/moleculare
-
-
-PopulationsLiteral = Literal["STIM", "PYR", "PVBC", "OLM"]
-
-
 def AllowStringsFrom(enum):
     """For convenience, allows users to specify enum values using their string name"""
 
@@ -118,13 +93,10 @@ def _cast(v) -> int:
     SynapseTypesDef, AllowStringsFrom(SynapseTypesDef)
 ]
 SWCTypesDefOrStr = Annotated[SWCTypesDef, AllowStringsFrom(SWCTypesDef)]
-LayersDefOrStr = Annotated[LayersDef, AllowStringsFrom(LayersDef)]
 SynapseMechanismsDefOrStr = Annotated[
     SynapseMechanismsDef, AllowStringsFrom(SynapseMechanismsDef)
 ]
-PopulationsDefOrStr = Annotated[
-    PopulationsDef, AllowStringsFrom(PopulationsDef)
-]
+LayersDefOrStr = Annotated[LayersDef, AllowStringsFrom(LayersDef)]
 
 
 PopulationName = str
@@ -253,7 +225,7 @@ class CellType(BaseModel):
             Dict[
                 SynapseTypesLiteral,
                 Dict[
-                    LayersLiteral,
+                    str,
                     Dict[Literal["mean", "variance"], float],
                 ],
             ],
@@ -262,15 +234,15 @@ class CellType(BaseModel):
     mechanism: Optional[Dict] = None
 
 
-CellTypes = Dict[PopulationsLiteral, CellType]
+CellTypes = Dict[str, CellType]
 
 
 class AxonExtent(BaseModel):
     width: Tuple[float, float]
     offset: Tuple[float, float]
 
 
-AxonExtents = Dict[PopulationsLiteral, Dict[LayerName, AxonExtent]]
+AxonExtents = Dict[str, Dict[LayerName, AxonExtent]]
 
 
 def probabilities_sum_to_one(x):

diff --git a/src/miv_simulator/simulator/generate_connections.py b/src/miv_simulator/simulator/generate_connections.py
@@ -91,6 +91,7 @@ def generate_distance_connections(
         coordinates_namespace=coords_namespace,
         synapses_namespace=synapses_namespace,
         distances_namespace=distances_namespace,
+        populations_dict=env.Populations,
         io_size=io_size,
         chunk_size=chunk_size,
         value_chunk_size=value_chunk_size,
@@ -112,6 +113,7 @@ def generate_connections(
     coordinates_namespace: str,
     synapses_namespace: str,
     distances_namespace: str,
+    populations_dict: Dict[str, int],
     io_size: int,
     chunk_size: int,
     value_chunk_size: int,
@@ -236,7 +238,6 @@ def generate_connections(
             r = random.Random(seeds)
             seeds = [r.randint(0, 2**32 - 1) for _ in range(3)]
 
-        populations_dict = config.PopulationsDef.__members__
         generate_uv_distance_connections(
             comm,
             populations_dict,