BinaryNetwork.ClusteredEI_network

  1"""Clustered E/I model built on top of the generic binary-network engine."""
  2
  3from __future__ import annotations
  4
  5import math
  6from typing import Callable, Dict, List, Optional, Sequence, Tuple
  7
  8import numpy as np
  9
 10from .BinaryNetwork import (
 11    AllToAllSynapse,
 12    BackgroundActivity,
 13    BinaryNetwork as BaseBinaryNetwork,
 14    BinaryNeuronPopulation,
 15    FixedIndegreeSynapse,
 16    Neuron,
 17    PairwiseBernoulliSynapse,
 18    PoissonSynapse,
 19)
 20
 21
 22def _total_population(parameter: Dict) -> float:
 23    if "N_E" not in parameter or "N_I" not in parameter:
 24        raise ValueError("Parameters must define both N_E and N_I.")
 25    return float(parameter["N_E"]) + float(parameter["N_I"])
 26
 27
 28def _normalize_conn_type(kind: Optional[str]) -> str:
 29    label = "bernoulli" if kind is None else str(kind).replace("_", "-").lower()
 30    if label not in {"bernoulli", "poisson", "fixed-indegree"}:
 31        raise ValueError(f"Unknown connection_type '{kind}'. Expected 'bernoulli', 'poisson', or 'fixed-indegree'.")
 32    return label
 33
 34
 35def _split_counts(total: int, groups: int) -> List[int]:
 36    if groups <= 0:
 37        raise ValueError("Q must be positive.")
 38    if total % groups != 0:
 39        raise ValueError(f"Population size {total} must be divisible by Q={groups}.")
 40    return [total // groups] * groups
 41
 42
 43def _mix_scales(R_plus: float, Q: int, kappa: float) -> Tuple[float, float, float, float]:
 44    positive = max(float(R_plus), 0.0)
 45    if Q <= 1:
 46        value = positive
 47        weight = value ** kappa
 48        return value ** (1.0 - kappa), value ** (1.0 - kappa), weight, weight
 49    prob_in = positive ** (1.0 - kappa)
 50    prob_out = (Q - prob_in) / (Q - 1)
 51    weight_in = positive ** kappa
 52    weight_out = (Q - weight_in) / (Q - 1)
 53    return prob_in, prob_out, weight_in, weight_out
 54
 55
 56def _compute_cluster_parameters(parameter: Dict, kappa: float) -> Dict[str, np.ndarray]:
 57    Q = int(parameter["Q"])
 58    N_E = float(parameter["N_E"])
 59    N_I = float(parameter["N_I"])
 60    N = _total_population(parameter)
 61    V_th = float(parameter["V_th"])
 62    g = float(parameter["g"])
 63    p0_ee = float(parameter["p0_ee"])
 64    p0_ie = float(parameter["p0_ie"])
 65    p0_ei = float(parameter["p0_ei"])
 66    p0_ii = float(parameter["p0_ii"])
 67    R_Eplus = parameter.get("R_Eplus")
 68    if R_Eplus is None:
 69        raise ValueError("R_Eplus must be set for binary simulations.")
 70    R_Eplus = float(R_Eplus)
 71    R_j = float(parameter.get("R_j", 0.0))
 72
 73    n_er = N_E / N
 74    n_ir = N_I / N
 75    theta_E = V_th
 76    theta_I = V_th
 77    R_Iplus = 1.0 + R_j * (R_Eplus - 1.0)
 78
 79    j_EE = theta_E / math.sqrt(p0_ee * n_er)
 80    j_IE = theta_I / math.sqrt(p0_ie * n_er)
 81    j_EI = -g * j_EE * p0_ee * n_er / (p0_ei * n_ir)
 82    j_II = -j_IE * p0_ie * n_er / (p0_ii * n_ir)
 83
 84    scale = 1.0 / math.sqrt(N)
 85    j_EE *= scale
 86    j_IE *= scale
 87    j_EI *= scale
 88    j_II *= scale
 89
 90    prob_in_E, prob_out_E, weight_in_E, weight_out_E = _mix_scales(R_Eplus, Q, kappa)
 91    prob_in_I, prob_out_I, weight_in_I, weight_out_I = _mix_scales(R_Iplus, Q, kappa)
 92
 93    P_EE = p0_ee * prob_in_E
 94    p_ee = p0_ee * prob_out_E
 95    P_EI = p0_ei * prob_in_I
 96    p_ei = p0_ei * prob_out_I
 97    P_IE = p0_ie * prob_in_I
 98    p_ie = p0_ie * prob_out_I
 99    P_II = p0_ii * prob_in_I
100    p_ii = p0_ii * prob_out_I
101
102    J_EE = j_EE * weight_in_E
103    j_ee = j_EE * weight_out_E
104    J_EI = j_EI * weight_in_I
105    j_ei = j_EI * weight_out_I
106    J_IE = j_IE * weight_in_I
107    j_ie = j_IE * weight_out_I
108    J_II = j_II * weight_in_I
109    j_ii = j_II * weight_out_I
110
111    J_EX = math.sqrt(p0_ee * N_E)
112    J_IX = 0.8 * J_EX
113
114    return {
115        "p_plus": np.array([[P_EE, P_EI], [P_IE, P_II]], dtype=float),
116        "p_minus": np.array([[p_ee, p_ei], [p_ie, p_ii]], dtype=float),
117        "j_plus": np.array([[J_EE, J_EI], [J_IE, J_II]], dtype=float),
118        "j_minus": np.array([[j_ee, j_ei], [j_ie, j_ii]], dtype=float),
119        "theta_E": theta_E,
120        "theta_I": theta_I,
121        "external_exc": J_EX,
122        "external_inh": J_IX,
123    }
124
125
126def _flatten_values(values) -> List:
127    if values is None:
128        return [None]
129    if isinstance(values, np.ndarray):
130        return values.flatten().tolist()
131    if isinstance(values, (list, tuple)):
132        return list(values)
133    return [values]
134
135
136def _normalize_activity_entry(entry, count: int, default_mode: str) -> List[Optional[Dict[str, float]]]:
137    if entry is None:
138        return [None] * count
139    mode = default_mode
140    values = entry
141    if isinstance(entry, dict):
142        mode = entry.get("mode", default_mode)
143        if "values" in entry:
144            values = entry["values"]
145        elif "value" in entry:
146            values = entry["value"]
147        else:
148            values = None
149    items = _flatten_values(values)
150    if len(items) == 1 and count > 1:
151        items = items * count
152    if len(items) != count:
153        raise ValueError(f"Initializer definition must provide {count} entries, got {len(items)}.")
154    normalized: List[Optional[Dict[str, float]]] = []
155    for value in items:
156        if value is None:
157            normalized.append(None)
158        else:
159            normalized.append({"mode": mode, "value": float(value)})
160    return normalized
161
162
163def _make_initializer(spec: Optional[Dict[str, float]], size: int) -> Optional[Callable[[int], np.ndarray]]:
164    if spec is None:
165        return None
166    mode = spec.get("mode", "bernoulli").lower()
167    value = float(spec.get("value", 0.0))
168    value = min(max(value, 0.0), 1.0)
169    if mode == "deterministic":
170        def initializer(_count=size, frac=value):
171            ones = int(round(frac * _count))
172            ones = min(max(ones, 0), _count)
173            state = np.zeros(_count, dtype=np.int16)
174            if ones > 0:
175                state[:ones] = 1
176                np.random.shuffle(state)
177            return state
178        return initializer
179    if mode == "bernoulli":
180        def initializer(_count=size, prob=value):
181            return (np.random.random(_count) < prob).astype(np.int16)
182        return initializer
183    raise ValueError(f"Unsupported initializer mode '{mode}'.")
184
185
186def _resolve_initializers(config, excit_sizes: Sequence[int], inhib_sizes: Sequence[int]) -> Tuple[List, List]:
187    if config is None:
188        return [None] * len(excit_sizes), [None] * len(inhib_sizes)
189    if not isinstance(config, dict):
190        config = {"default": config}
191    default_mode = str(config.get("mode", "bernoulli")).lower()
192    excit_specs: List[Optional[Dict[str, float]]] = [None] * len(excit_sizes)
193    inhib_specs: List[Optional[Dict[str, float]]] = [None] * len(inhib_sizes)
194
195    def apply(entry, targets):
196        if entry is None:
197            return
198        normalized = _normalize_activity_entry(entry, len(targets), default_mode)
199        for idx, spec in enumerate(normalized):
200            if spec is not None:
201                targets[idx] = spec
202
203    apply(config.get("default"), excit_specs)
204    apply(config.get("default"), inhib_specs)
205    apply(config.get("excitatory"), excit_specs)
206    apply(config.get("inhibitory"), inhib_specs)
207    apply(config.get("excitatory_by_cluster"), excit_specs)
208    apply(config.get("inhibitory_by_cluster"), inhib_specs)
209
210    excit_initializers = [_make_initializer(spec, size) for spec, size in zip(excit_specs, excit_sizes)]
211    inhib_initializers = [_make_initializer(spec, size) for spec, size in zip(inhib_specs, inhib_sizes)]
212    return excit_initializers, inhib_initializers
213
214
215class ClusteredEI_network(BaseBinaryNetwork):
216    """Repository-specific clustered E/I binary network.
217
218    Examples
219    --------
220    >>> parameter = {
221    ...     "Q": 2,
222    ...     "N_E": 4,
223    ...     "N_I": 2,
224    ...     "V_th": 1.0,
225    ...     "g": 1.0,
226    ...     "p0_ee": 0.2,
227    ...     "p0_ei": 0.2,
228    ...     "p0_ie": 0.2,
229    ...     "p0_ii": 0.2,
230    ...     "R_Eplus": 1.0,
231    ...     "R_j": 0.0,
232    ...     "m_X": 0.0,
233    ...     "tau_e": 5.0,
234    ...     "tau_i": 10.0,
235    ... }
236    >>> net = ClusteredEI_network(parameter)
237    >>> net.Q
238    2
239    """
240
241    def __init__(self, parameter: Dict, *, kappa: Optional[float] = None, connection_type: Optional[str] = None, name="Binary EI Network"):
242        super().__init__(name)
243        self.parameter = dict(parameter)
244        self.Q = int(self.parameter["Q"])
245        self.connection_type = _normalize_conn_type(connection_type or self.parameter.get("connection_type"))
246        self.multapses = bool(self.parameter.get("multapses", True))
247        self.kappa = float(kappa if kappa is not None else self.parameter.get("kappa", 0.0))
248        self.connection_parameters = _compute_cluster_parameters(self.parameter, self.kappa)
249        self.E_sizes = _split_counts(int(self.parameter["N_E"]), self.Q)
250        self.I_sizes = _split_counts(int(self.parameter["N_I"]), self.Q)
251        self.total_neurons = sum(self.E_sizes) + sum(self.I_sizes)
252        self.initializers_E, self.initializers_I = _resolve_initializers(
253            self.parameter.get("initial_activity"),
254            self.E_sizes,
255            self.I_sizes,
256        )
257        self.E_pops: List[BinaryNeuronPopulation] = []
258        self.I_pops: List[BinaryNeuronPopulation] = []
259        self.other_pops: List[Neuron] = []
260        self._structure_created = False
261
262    def _build_populations(self):
263        tau_e = float(self.parameter["tau_e"])
264        tau_i = float(self.parameter.get("tau_i", 0.5 * tau_e))
265        theta_E = self.connection_parameters["theta_E"]
266        theta_I = self.connection_parameters["theta_I"]
267        for idx, size in enumerate(self.E_sizes):
268            pop = BinaryNeuronPopulation(
269                self,
270                N=size,
271                threshold=theta_E,
272                tau=tau_e,
273                name=f"E{idx}",
274                initializer=self.initializers_E[idx],
275            )
276            pop.cluster_index = idx
277            pop.cell_type = "E"
278            self.E_pops.append(self.add_population(pop))
279        for idx, size in enumerate(self.I_sizes):
280            pop = BinaryNeuronPopulation(
281                self,
282                N=size,
283                threshold=theta_I,
284                tau=tau_i,
285                name=f"I{idx}",
286                initializer=self.initializers_I[idx],
287            )
288            pop.cluster_index = idx
289            pop.cell_type = "I"
290            self.I_pops.append(self.add_population(pop))
291
292    def _synapse_factory(self):
293        if self.connection_type == "poisson":
294            return lambda pre, post, p, j: PoissonSynapse(self, pre, post, rate=p, j=j)
295        if self.connection_type == "fixed-indegree":
296            return lambda pre, post, p, j: FixedIndegreeSynapse(
297                self,
298                pre,
299                post,
300                p=p,
301                j=j,
302                multapses=self.multapses,
303            )
304        return lambda pre, post, p, j: PairwiseBernoulliSynapse(self, pre, post, p=p, j=j)
305
306    def _build_synapses(self):
307        builder = self._synapse_factory()
308        p_plus = self.connection_parameters["p_plus"]
309        p_minus = self.connection_parameters["p_minus"]
310        j_plus = self.connection_parameters["j_plus"]
311        j_minus = self.connection_parameters["j_minus"]
312        pre_groups = [self.E_pops, self.I_pops]
313        for target_group_idx, target_pops in enumerate([self.E_pops, self.I_pops]):
314            for source_group_idx, source_pops in enumerate(pre_groups):
315                for post_pop in target_pops:
316                    for pre_pop in source_pops:
317                        same_cluster = pre_pop.cluster_index == post_pop.cluster_index
318                        p_matrix = p_plus if same_cluster else p_minus
319                        j_matrix = j_plus if same_cluster else j_minus
320                        p_value = p_matrix[target_group_idx, source_group_idx]
321                        j_value = j_matrix[target_group_idx, source_group_idx]
322                        self.add_synapse(builder(pre_pop, post_pop, p_value, j_value))
323
324    def _build_background(self):
325        m_X = float(self.parameter.get("m_X", 0.0) or 0.0)
326        if m_X == 0.0:
327            return
328        bg = BackgroundActivity(self, N=1, Activity=m_X, name="Background")
329        #self.other_pops.append(self.add_population(bg))
330        J_EX = float(self.connection_parameters["external_exc"])
331        J_IX = float(self.connection_parameters["external_inh"])
332        for pop in self.E_pops:
333            pop.threshold-=m_X*J_EX
334            #self.add_synapse(AllToAllSynapse(self, bg, pop, j=J_EX))
335        for pop in self.I_pops:
336            pop.threshold -= m_X * J_IX
337            #self.add_synapse(AllToAllSynapse(self, bg, pop, j=J_IX))
338
339    def _ensure_structure(self):
340        if self._structure_created:
341            return
342        self._build_populations()
343        self._build_synapses()
344        self._build_background()
345        self._structure_created = True
346
347    def initialize(
348        self,
349        autapse: bool = False,
350        weight_mode: str = "auto",
351        ram_budget_gb: float = 12.0,
352        weight_dtype=np.float32,
353    ):
354        """Build clustered populations/synapses and initialize the parent network.
355
356        Expected output
357        ---------------
358        After initialization, `state`, `field`, and the sampled connectivity are
359        allocated and ready for `run(...)`.
360        """
361        self._ensure_structure()
362        super().initialize(
363            autapse=autapse,
364            weight_mode=weight_mode,
365            ram_budget_gb=ram_budget_gb,
366            weight_dtype=weight_dtype,
367        )
368
369    def reinitalize(self):
370        """Rebuild the simulation state with the stored parameters."""
371        self.initialize()