More type checker fixes

GazzolaLab · Apr 6, 2024 · 6204bec · 6204bec
1 parent f2402b5
commit 6204bec
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/src/miv_simulator/config.py b/src/miv_simulator/config.py
@@ -119,8 +119,8 @@ def __getitem__(self, item):
 
 
 class Mechanism(BaseModel):
-    g_unit: float
-    weight: float
+    g_unit: Optional[float] = None
+    weight: Optional[float] = None
     tau_rise: Optional[float] = None
     tau_decay: Optional[float] = None
     e: Optional[int] = None
@@ -131,7 +131,7 @@ class Synapse(BaseModel):
     sections: List[SWCTypesLiteral]
     layers: List[LayerName]
     proportions: list[float]
-    mechanisms: Dict[SynapseMechanismName, Mechanism]
+    mechanisms: Dict[SynapseMechanismName, Dict[Union[str, int], Mechanism]]
     contacts: int = 1
 
     def to_config(self, layer_definitions: Dict[LayerName, int]):
@@ -314,7 +314,21 @@ def axon_extents(self) -> AxonExtents:
 
     @property
     def synapses(self) -> Synapses:
-        return self.get("Connection Generator.Synapses")
+        synapses = {}
+        for post, v in self.get("Connection Generator.Synapses").items():
+            synapses[post] = {}
+            for pre, syn in v.items():
+                synapse = copy.copy(syn)
+                if "swctype mechanisms" in synapse:
+                    synapse["mechanisms"] = {
+                        SWCTypesDef.__members__[swc_type]: c
+                        for swc_type, c in synapse["swctype mechanisms"].items()
+                    }
+                    del synapse["swctype mechanisms"]
+                elif "mechanisms" in synapse:
+                    synapse["mechanisms"] = {"default": synapse["mechanisms"]}
+                synapses[post][pre] = synapse
+        return synapses
 
     @property
     def projections(self) -> SynapticProjections:

diff --git a/src/miv_simulator/utils/io.py b/src/miv_simulator/utils/io.py
@@ -274,7 +274,7 @@ def create_neural_h5(
 ) -> None:
     populations = []
     for pop_name, pop_idx in population_definitions.items():
-        layer_counts = cell_distributions[pop_name]
+        layer_counts = cell_distributions.get(pop_name, {})
         pop_count = 0
         for layer_name, layer_count in layer_counts.items():
             pop_count += layer_count