enable cell_permute for coreneuron

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=[unpermuted, node-adjacency]   Cell permutation [default: unpermuted].
+                                Only available for CoreNEURON.
+                                Currently incompatible with NEURON. Options:
+                                - unpermuted: No permutation
+                                - node-adjacency: 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):
+    UNPERMUTED = 0  # cpu
+    NODE_ADJACENCY = 1  # cpu/gpu
+
+    __mapping__ = [
+        ("unpermuted", UNPERMUTED),
+        ("node-adjacency", NODE_ADJACENCY),
+    ]
+    __default__ = UNPERMUTED
+
+
 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,18 @@ 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.UNPERMUTED:
+        raise ConfigurationError(
+            f"Cell permutation is only available with CoreNEURON. "
+            f"--cell-permute={config.cell_permute.to_string()} is invalid with NEURON."
+        )
+
+
 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/integration-e2e/test_reports.py

@@ -11,6 +11,7 @@
 from ..utils import ReportReader
 import copy
 from neurodamus.utils.pyutils import CumulativeError
+from neurodamus.core.coreneuron_simulation_config import CoreSimulationConfig
 
 
 _BASE_EXTRA_CONFIG = {
@@ -406,6 +407,7 @@ def test_compartment_missing_ref(create_tmp_simulation_config_file):
                 },
             },
         }
+
     ],
     indirect=True,
 )
@@ -430,3 +432,36 @@ def test_results_are_identical_with_single_report(create_tmp_simulation_config_f
     r_reference = ReportReader(reference_dir / file_name)   
     r = ReportReader(output_dir / file_name )
     assert r_reference == r
+
+
+@pytest.mark.parametrize(
+    "create_tmp_simulation_config_file",
+    [
+        make_extra_config("v5_sonata_config", "CORENEURON"),
+    ],
+    indirect=True,
+)
+def test_reports_cell_permute(create_tmp_simulation_config_file):
+    """
+    Test that enabling cell permutation (cell_permute=node-adjacency) preserves report consistency.
+    """
+    nd = Neurodamus(create_tmp_simulation_config_file, cell_permute="node-adjacency", keep_build=True)
+    output_dir = Path(SimConfig.output_root)
+    reference_dir = V5_SONATA / "reference" / "reports"
+    sim_conf = CoreSimulationConfig.load("build/sim.conf")
+    assert sim_conf.cell_permute == 1
+
+    nd.run()
+    loose_tols = {"rtol": 1e-6, "atol": 1e-6}
+
+    # Compare files to reference. Since the reference is fixed, this is also a comparison neuron vs coreneuron
+    # reference produced with neuron
+    # coreneuron does not have exactly the same results, we use the loose tols in that case
+    loose_tol_files = {"summation_i_membrane.h5"}
+    for ref_file in reference_dir.glob("*.h5"):
+        r_reference = ReportReader(ref_file)   
+        file = output_dir / ref_file.name 
+        r = ReportReader(file)
+
+        assert r.allclose(r_reference, **(loose_tols if ref_file.name in loose_tol_files else {})), f"The reports differ:\n{file}\n{ref_file}"
+

tests/unit/test_cell_permute.py

@@ -0,0 +1,40 @@
+import pytest
+
+from neurodamus.core.configuration import SimConfig
+from neurodamus.core.configuration import CellPermute
+
+from neurodamus.core.coreneuron_simulation_config import CoreSimulationConfig
+
+
+@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):
+    from neurodamus import Neurodamus
+    Neurodamus(create_tmp_simulation_config_file, cell_permute="node-adjacency", keep_build=True)
+    assert SimConfig.cell_permute == CellPermute.NODE_ADJACENCY
+    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):
+    from neurodamus import Neurodamus
+    Neurodamus(create_tmp_simulation_config_file, keep_build=True)
+    assert SimConfig.cell_permute == CellPermute.UNPERMUTED
+    sim_conf = CoreSimulationConfig.load("build/sim.conf")
+    assert sim_conf.cell_permute == 0
+

tests/unit/test_cli.py

@@ -5,6 +5,7 @@
 import pytest
 
 from neurodamus.core.coreneuron_report_config import CoreReportConfig
+from neurodamus.core.coreneuron_simulation_config import CoreSimulationConfig
 
 @pytest.mark.parametrize(
     "create_tmp_simulation_config_file",
@@ -196,3 +197,47 @@ 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=node-adjacency", "--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_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