WIP

Author: cattabiani

neurodamus/commands.py

@@ -71,6 +71,11 @@ def neurodamus(args=None):
         --report-buffer-size=<number> Override the size in MB each rank will allocate for each
                                       report buffer to hold data. When the buffer is full, the
                                       ranks will aggregate data for writing to disk. Default: 8 MB
+        --cell-permute=[0, 1]   Cell permutation [default: 0].
+                                Only available for CoreNEURON.
+                                Currently incompatible with NEURON. Options:
+                                - 0  No permutation
+                                - 1  Optimise for node adjacency
     """
     from . import __version__
 

neurodamus/core/configuration.py

@@ -11,7 +11,7 @@
 from ._shmutils import SHMUtil
 from neurodamus.io.sonata_config import SonataConfig
 from neurodamus.utils.logging import log_verbose
-from neurodamus.utils.pyutils import ConfigT
+from neurodamus.utils.pyutils import ConfigT, StrEnumBase
 
 EXCEPTION_NODE_FILENAME = ".exception_node"
 """A file which controls which rank shows exception"""
@@ -54,7 +54,19 @@ class Feature(Enum):
     LoadBalance = 5
 
 
+class CellPermute(StrEnumBase):
+    NONE = 0  # cpu
+    NODE_ADJACENCY = 1  # cpu/gpu
+
+    __mapping__ = [
+        ("0", NONE),
+        ("1", NODE_ADJACENCY),
+    ]
+    __default__ = NONE
+
+
 class CliOptions(ConfigT):
+    cell_permute = None
     report_buffer_size = None
     build_model = None
     simulate_model = True
@@ -205,6 +217,7 @@ class _SimConfig:
     use_coreneuron = False
     use_neuron = True
     report_buffer_size = 8  # in MB
+    cell_permute = CellPermute.default()
     delete_corenrn_data = False
     modelbuilding_steps = 1
     build_model = True
@@ -1052,6 +1065,19 @@ def _coreneuron_direct_mode(config: _SimConfig):
     config.coreneuron_direct_mode = direct_mode
 
 
+@SimConfig.validator
+def _cell_permute(config: _SimConfig):
+    user_config = config.cli_options
+    if user_config.cell_permute is not None:
+        config.cell_permute = CellPermute.from_string(str(user_config.cell_permute))
+    if config.use_neuron and config.cell_permute != CellPermute.NONE:
+        logging.warning(
+            "Cell permutation is only available with CoreNEURON. "
+            "--cell-permute=%s will be ignored.",
+            config.cell_permute.to_string(),
+        )
+
+
 def get_debug_cell_gids(cli_options):
     """Parse the --dump-cell-state option from CLI.
 

neurodamus/node.py

@@ -1287,7 +1287,7 @@ def _coreneuron_write_sim_config(self, corenrn_restore):
             prcellgid=prcellgid,
             celsius=getattr(SimConfig, "celsius", 34.0),
             voltage=getattr(SimConfig, "v_init", -65.0),
-            cell_permute=CoreConfig.default_cell_permute,
+            cell_permute=int(SimConfig.cell_permute),
             pattern=self._core_replay_file or None,
             seed=SimConfig.rng_info.getGlobalSeed(),
             model_stats=int(SimConfig.cli_options.model_stats),

neurodamus/report_parameters.py

@@ -2,11 +2,10 @@
 
 import logging
 from dataclasses import dataclass
-from enum import IntEnum
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from .utils.pyutils import cache_errors
+from .utils.pyutils import StrEnumBase, cache_errors
 
 if TYPE_CHECKING:
     from neurodamus.target_manager import TPointList
@@ -16,28 +15,6 @@ class ReportSetupError(Exception):
     pass
 
 
-class StrEnumBase(IntEnum):
-    __mapping__: list[tuple[str, int]] = []
-    # default for when there is value. Leaving None throws an error
-    __default__ = None
-    # default when the string is not found in the mapping
-    __invalid__ = None
-
-    @classmethod
-    def from_string(cls, s: str):
-        if not s:
-            return cls(cls.__default__)
-        mapping = dict(cls.__mapping__)
-        return cls(mapping.get(s.lower(), cls.__invalid__))
-
-    def to_string(self) -> str:
-        reverse__mapping__ = {v: k for k, v in self.__mapping__}
-        return reverse__mapping__[self]
-
-    def __str__(self):
-        return f"{self.__class__.__name__}.{self.name}"
-
-
 class SectionType(StrEnumBase):
     ALL = 0
     SOMA = 1

neurodamus/utils/pyutils.py

@@ -5,11 +5,37 @@
 import subprocess  # noqa: S404
 import weakref
 from bisect import bisect_left
-from enum import EnumMeta
+from enum import EnumMeta, IntEnum
 
 import numpy as np
 
 
+class StrEnumBase(IntEnum):
+    __mapping__: list[tuple[str, int]] = []
+    # default for when there is value. Leaving None throws an error
+    __default__ = None
+    # default when the string is not found in the mapping
+    __invalid__ = None
+
+    @classmethod
+    def from_string(cls, s: str):
+        if not s:
+            return cls(cls.__default__)
+        mapping = dict(cls.__mapping__)
+        return cls(mapping.get(s.lower(), cls.__invalid__ if cls.__invalid__ is not None else s))
+
+    def to_string(self) -> str:
+        reverse__mapping__ = {v: k for k, v in self.__mapping__}
+        return reverse__mapping__[self]
+
+    def __str__(self):
+        return f"{self.__class__.__name__}.{self.name}"
+
+    @classmethod
+    def default(cls):
+        return cls(cls.__default__)
+
+
 class CumulativeError(Exception):
     def __init__(self, errors=None):
         self.errors = errors or []

tests/unit/test_cli.py

@@ -5,6 +5,9 @@
 import pytest
 
 from neurodamus.core.coreneuron_report_config import CoreReportConfig
+from neurodamus.core.coreneuron_simulation_config import CoreSimulationConfig
+from neurodamus.core.configuration import SimConfig
+from neurodamus.core.configuration import CellPermute
 
 @pytest.mark.parametrize(
     "create_tmp_simulation_config_file",
@@ -196,3 +199,75 @@ def test_cli_report_buff_invalid(create_tmp_simulation_config_file):
         result = subprocess.run(command, check=False, capture_output=True, text=True)
 
         assert "Report buffer size must be > 0" in result.stdout
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [{"simconfig_fixture": "ringtest_baseconfig",
+              "extra_config": {
+            "target_simulator": "CORENEURON",
+        }
+      }],
+    indirect=True,
+)
+def test_cli_cell_permute_simple_setting(create_tmp_simulation_config_file):
+    command = ["neurodamus", create_tmp_simulation_config_file, "--cell-permute=1", "--keep-build"]
+    subprocess.run(command, check=False, capture_output=True, text=True)
+    sim_conf = CoreSimulationConfig.load("build/sim.conf")
+    assert sim_conf.cell_permute == 1
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [{"simconfig_fixture": "ringtest_baseconfig",
+              "extra_config": {
+            "target_simulator": "CORENEURON",
+        }
+      }],
+    indirect=True,
+)
+def test_cli_cell_permute_default(create_tmp_simulation_config_file):
+    command = ["neurodamus", create_tmp_simulation_config_file, "--keep-build"]
+    subprocess.run(command, check=False, capture_output=True, text=True)
+    sim_conf = CoreSimulationConfig.load("build/sim.conf")
+    assert sim_conf.cell_permute == 0
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [{"simconfig_fixture": "ringtest_baseconfig",
+              "extra_config": {
+            "target_simulator": "CORENEURON",
+        }
+      }],
+    indirect=True,
+)
+def test_cli_cell_permute_simple_setting_python(create_tmp_simulation_config_file):
+    from neurodamus import Neurodamus
+    nd = Neurodamus(create_tmp_simulation_config_file, cell_permute=1)
+    assert SimConfig.cell_permute == CellPermute.NODE_ADJACENCY
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [{"simconfig_fixture": "ringtest_baseconfig",
+              "extra_config": {
+            "target_simulator": "CORENEURON",
+        }
+      }],
+    indirect=True,
+)
+def test_cli_cell_permute_default_python(create_tmp_simulation_config_file):
+    from neurodamus import Neurodamus
+    nd = Neurodamus(create_tmp_simulation_config_file)
+    assert SimConfig.cell_permute == CellPermute.NONE
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [{"simconfig_fixture": "ringtest_baseconfig",
+              "extra_config": {
+            "target_simulator": "CORENEURON",
+        }
+      }],
+    indirect=True,
+)            
+def test_cli_cell_permute_invalid(create_tmp_simulation_config_file):
+    command = ["neurodamus", create_tmp_simulation_config_file, "--cell-permute=2"]
+    result = subprocess.run(command, check=False, capture_output=True, text=True)
+    assert "'2' is not a valid CellPermute" in result.stdout