sim_config

  1"""Hierarchical configuration loading, merging, and override helpers.
  2
  3The functions in this module are the common configuration entry point for the
  4figure scripts and pipelines.
  5
  6Simple hierarchical YAML example:
  7
  8```yaml
  9# my_config.yaml
 10Q: 4
 11R_Eplus: 6.0
 12binary:
 13  seed: 3
 14  simulation_steps: 200000
 15  sample_interval: 10
 16plot:
 17  raster:
 18    stride: 4
 19```
 20
 21Loading that config in another script:
 22
 23```python
 24from sim_config import load_config
 25
 26cfg = load_config("my_config.yaml")
 27binary_seed = cfg["binary"]["seed"]
 28```
 29
 30Applying overrides the same way the figure scripts do:
 31
 32```python
 33from sim_config import load_config
 34
 35cfg = load_config(
 36    "my_config.yaml",
 37    overrides=[
 38        "binary.seed=7",
 39        "binary.simulation_steps=500000",
 40        "plot.raster.stride=2",
 41    ],
 42)
 43```
 44
 45Using the argparse helpers:
 46
 47```python
 48import argparse
 49
 50from sim_config import add_override_arguments, load_from_args
 51
 52parser = argparse.ArgumentParser()
 53add_override_arguments(parser)
 54args = parser.parse_args()
 55cfg = load_from_args(args)
 56```
 57
 58What this looks like on the command line:
 59
 60```python
 61# script.py
 62import argparse
 63
 64from sim_config import add_override_arguments, load_from_args
 65
 66parser = argparse.ArgumentParser()
 67add_override_arguments(parser)
 68args = parser.parse_args()
 69cfg = load_from_args(args)
 70print(cfg["binary"]["seed"], cfg["plot"]["raster"]["stride"])
 71```
 72
 73```bash
 74python script.py --config my_config.yaml \
 75  -O binary.seed=7 \
 76  -O plot.raster.stride=2
 77```
 78
 79Expected output
 80---------------
 81
 82`cfg["binary"]["seed"] == 7` and `cfg["plot"]["raster"]["stride"] == 2`.
 83
 84Using `sim_tag_from_cfg(...)` to run a parameter sweep and save one folder per
 85configuration:
 86
 87```python
 88from pathlib import Path
 89
 90from sim_config import load_config, sim_tag_from_cfg, write_yaml_config
 91
 92
 93def mock_run(cfg):
 94    return {
 95        "rate_hz": cfg["R_Eplus"] * 1.5,
 96        "seed": cfg["binary"]["seed"],
 97    }
 98
 99
100output_root = Path("mock_results")
101
102for q in [2, 4, 6]:
103    cfg = load_config(
104        "default_simulation",
105        overrides=[
106            f"Q={q}",
107            f"binary.seed={100 + q}",
108        ],
109    )
110    tag = sim_tag_from_cfg(cfg)
111    run_dir = output_root / f"Q{q}" / tag
112    run_dir.mkdir(parents=True, exist_ok=True)
113
114    result = mock_run(cfg)
115    write_yaml_config(cfg, run_dir / "config.yaml")
116    (run_dir / "result.txt").write_text(f"{result}\n", encoding="utf-8")
117```
118
119This creates a deterministic folder structure such as
120`mock_results/Q4/3f7a1c2b9e/`, writes the exact configuration snapshot next to
121the result, and keeps reruns with identical settings in the same tagged
122location.
123
124Regenerating docs:
125
126```bash
127python scripts/generate_api_docs.py
128```
129"""
130
131from __future__ import annotations
132
133import argparse
134import hashlib
135import json
136import math
137from copy import deepcopy
138from pathlib import Path
139from typing import Any, Iterable, Mapping, Optional, Sequence, Union
140
141import numpy as np
142
143try:
144    import yaml  # type: ignore
145except ModuleNotFoundError as _yaml_error:  # pragma: no cover - optional dependency
146    yaml = None  # type: ignore[assignment]
147else:
148    _yaml_error = None
149
150
151CONFIG_DIR = Path(__file__).resolve().parent
152
153__all__ = [
154    "write_yaml_config",
155    "sim_tag_from_cfg",
156    "load_config",
157    "parse_overrides",
158    "deep_update",
159    "add_override_arguments",
160    "load_from_args",
161]
162
163__pdoc__ = {
164    "CONFIG_DIR": False,
165    "_to_human": False,
166    "write_human_json": False,
167    "_normalize_for_tag": False,
168    "resolve_base_config": False,
169    "first_float": False,
170    "_resolve_config_path": False,
171}
172
173
174def _to_human(obj, *, float_precision=6, nan_policy="string"):
175    if isinstance(obj, np.ndarray):
176        return [
177            _to_human(x, float_precision=float_precision, nan_policy=nan_policy)
178            for x in obj.tolist()
179        ]
180    if isinstance(obj, (np.floating, np.integer, np.bool_)):
181        obj = obj.item()
182    if obj is np.nan:
183        return "NaN" if nan_policy == "string" else None
184    if obj is None or isinstance(obj, (bool, int, str)):
185        return obj
186    if isinstance(obj, float):
187        if math.isnan(obj):
188            return "NaN" if nan_policy == "string" else None
189        if math.isinf(obj):
190            return ("Infinity" if obj > 0 else "-Infinity") if nan_policy == "string" else None
191        return round(obj, float_precision) if float_precision is not None else obj
192    if isinstance(obj, (list, tuple)):
193        return [
194            _to_human(x, float_precision=float_precision, nan_policy=nan_policy)
195            for x in obj
196        ]
197    if isinstance(obj, dict):
198        return {
199            k: _to_human(v, float_precision=float_precision, nan_policy=nan_policy)
200            for k, v in obj.items()
201        }
202    return str(obj)
203
204
205def write_human_json(
206    cfg: dict,
207    path: Union[str, Path],
208    *,
209    float_precision: int = 6,
210    nan_policy: str = "string",
211    ensure_ascii: bool = False,
212):
213    path = Path(path)
214    path.parent.mkdir(parents=True, exist_ok=True)
215    human = _to_human(cfg, float_precision=float_precision, nan_policy=nan_policy)
216    txt = json.dumps(human, indent=2, sort_keys=True, ensure_ascii=ensure_ascii)
217    tmp = path.with_suffix(path.suffix + ".tmp")
218    tmp.write_text(txt, encoding="utf-8")
219    tmp.replace(path)
220
221def write_yaml_config(cfg: dict, path: Union[str, Path]) -> None:
222    """Write a configuration dictionary as YAML.
223
224    Examples
225    --------
226    >>> from pathlib import Path
227    >>> path = Path("/tmp/example_config.yaml")
228    >>> write_yaml_config({"binary": {"seed": 3}}, path)
229    >>> path.exists()
230    True
231    """
232    if yaml is None:  # pragma: no cover - optional dependency
233        raise ModuleNotFoundError(
234            "PyYAML is required to write configuration files. Install it via 'pip install pyyaml'."
235        ) from _yaml_error
236    path = Path(path)
237    path.parent.mkdir(parents=True, exist_ok=True)
238    with path.open("w", encoding="utf-8") as handle:
239        yaml.safe_dump(cfg, handle, sort_keys=True)
240
241
242def _normalize_for_tag(obj):
243    if obj is None or isinstance(obj, bool):
244        return obj
245    if isinstance(obj, int):
246        return int(obj)
247    if isinstance(obj, float):
248        if math.isnan(obj):
249            return {"__float__": "NaN"}
250        if math.isinf(obj):
251            return {"__float__": "Infinity" if obj > 0 else "-Infinity"}
252        return {"__float__": format(obj, ".17g")}
253    if isinstance(obj, (list, tuple)):
254        return [_normalize_for_tag(x) for x in obj]
255    if isinstance(obj, dict):
256        return {k: _normalize_for_tag(obj[k]) for k in sorted(obj.keys())}
257    return str(obj)
258
259
260def sim_tag_from_cfg(cfg: dict, *, length: int = 10) -> str:
261    """Create a stable short hash for a configuration dictionary.
262
263    Examples
264    --------
265    >>> sim_tag_from_cfg({"binary": {"seed": 3}}, length=8)
266    'da4c7220'
267    """
268    canon = _normalize_for_tag(cfg)
269    blob = json.dumps(canon, separators=(",", ":"), sort_keys=True, ensure_ascii=True)
270    return hashlib.sha1(blob.encode("utf-8")).hexdigest()[:length]
271
272
273def load_config(
274    name: str = "default_simulation",
275    *,
276    overrides: Optional[Iterable[str]] = None,
277) -> dict[str, Any]:
278    """Load a YAML configuration and optionally apply dotted overrides.
279
280    Examples
281    --------
282    ```python
283    cfg = load_config(
284        "default_simulation",
285        overrides=["binary.seed=7", "binary.simulation_steps=1000"],
286    )
287    ```
288
289    Expected output
290    ---------------
291    `cfg` is a nested dictionary where the override values replace the loaded
292    base configuration.
293    """
294    config_path = _resolve_config_path(name)
295    if yaml is None:  # pragma: no cover - optional dependency
296        raise ModuleNotFoundError(
297            "PyYAML is required to load configuration files. Install it via 'pip install pyyaml'."
298        ) from _yaml_error
299    with config_path.open(encoding="utf-8") as handle:
300        base_config = yaml.safe_load(handle) or {}
301    if not overrides:
302        return base_config
303    override_dict = parse_overrides(overrides)
304    return deep_update(base_config, override_dict)
305
306
307def resolve_base_config(descriptor: Any) -> dict[str, Any]:
308    if descriptor is None:
309        raise ValueError("Base configuration descriptor must not be None.")
310    if isinstance(descriptor, Mapping):
311        return deepcopy(descriptor)
312    if isinstance(descriptor, (str, Path)):
313        return load_config(str(descriptor))
314    raise TypeError(
315        "Base configuration must be provided as a mapping or a config name/path "
316        f"(got {type(descriptor).__name__})."
317    )
318
319
320def first_float(
321    value: Any,
322    *,
323    cell_type: Optional[str] = None,
324    default: Optional[float] = None,
325) -> float:
326    if value is None:
327        if default is None:
328            raise ValueError("No value available to coerce into float.")
329        return float(default)
330    if isinstance(value, (int, float)):
331        return float(value)
332    if isinstance(value, Mapping):
333        if cell_type and cell_type in value:
334            return first_float(value[cell_type], cell_type=cell_type, default=default)
335        if "default" in value:
336            return first_float(value["default"], cell_type=cell_type, default=default)
337        if "excitatory" in value:
338            return first_float(value["excitatory"], cell_type=cell_type, default=default)
339        first_key = next(iter(value))
340        return first_float(value[first_key], cell_type=cell_type, default=default)
341    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
342        if not value:
343            if default is None:
344                raise ValueError("Cannot extract float from empty sequence.")
345            return float(default)
346        return float(value[0])
347    raise TypeError(
348        "Expected a float-like value, sequence, or mapping; "
349        f"got {type(value).__name__}."
350    )
351
352
353def parse_overrides(pairs: Iterable[str]) -> dict[str, Any]:
354    """Parse CLI-style dotted overrides into a nested dictionary.
355
356    Examples
357    --------
358    >>> parse_overrides(["binary.seed=7", "plot.raster.stride=2"])
359    {'binary': {'seed': 7}, 'plot': {'raster': {'stride': 2}}}
360    """
361    root: dict[str, Any] = {}
362    for raw in pairs:
363        if "=" not in raw:
364            raise ValueError(f"Override '{raw}' is missing '='.")
365        key_path, raw_value = raw.split("=", 1)
366        target = root
367        *parents, leaf = key_path.split(".")
368        for segment in parents:
369            target = target.setdefault(segment, {})
370            if not isinstance(target, dict):
371                raise ValueError(f"Cannot override non-dict path '{key_path}'.")
372        target[leaf] = yaml.safe_load(raw_value)
373    return root
374
375
376def deep_update(base: dict[str, Any], updates: dict[str, Any]) -> dict[str, Any]:
377    """Recursively merge `updates` into a deep copy of `base`.
378
379    Examples
380    --------
381    >>> deep_update({"binary": {"seed": 3, "steps": 10}}, {"binary": {"steps": 20}})
382    {'binary': {'seed': 3, 'steps': 20}}
383    """
384    merged = deepcopy(base)
385    for key, value in updates.items():
386        if isinstance(value, dict) and isinstance(merged.get(key), dict):
387            merged[key] = deep_update(merged[key], value)
388        else:
389            merged[key] = deepcopy(value)
390    return merged
391
392
393def _resolve_config_path(name: str) -> Path:
394    candidate = Path(name)
395    if candidate.suffix:
396        return candidate if candidate.is_file() else CONFIG_DIR / candidate
397    return CONFIG_DIR / f"{name}.yaml"
398
399
400def add_override_arguments(
401    parser: argparse.ArgumentParser,
402    *,
403    config_option: str = "--config",
404    overwrite_option: str = "--overwrite",
405    config_default: str = "default_simulation",
406    overwrite_metavar: str = "path=value",
407) -> None:
408    """Add the repository's standard config/override arguments to an argparse parser.
409
410    Expected output
411    ---------------
412    The parser accepts a config argument and repeated `-O/--overwrite`
413    arguments compatible with the figure scripts.
414    """
415    parser.add_argument(
416        config_option,
417        default=config_default,
418        help="Config name or path (defaults to '%(default)s').",
419    )
420    parser.add_argument(
421        "-O",
422        overwrite_option,
423        action="append",
424        default=[],
425        metavar=overwrite_metavar,
426        help="Override a config value using dotted-path notation (may be repeated).",
427    )
428
429
430def load_from_args(
431    args: argparse.Namespace,
432    *,
433    config_attr: str = "config",
434    overwrite_attr: str = "overwrite",
435    default_config: str = "default_simulation",
436) -> dict[str, Any]:
437    """Load a configuration from parsed argparse arguments.
438
439    Examples
440    --------
441    ```python
442    parser = argparse.ArgumentParser()
443    add_override_arguments(parser)
444    args = parser.parse_args(["--config", "default_simulation", "-O", "binary.seed=9"])
445    cfg = load_from_args(args)
446    ```
447
448    Expected output
449    ---------------
450    `cfg["binary"]["seed"] == 9`.
451    """
452    config_name = getattr(args, config_attr, default_config) or default_config
453    overrides = getattr(args, overwrite_attr, None)
454    return load_config(config_name, overrides=overrides)