diff --git a/examples/nest-simulation/tests/test_examples.py b/examples/nest-simulation/tests/test_examples.py index f7c65953..28fe210f 100644 --- a/examples/nest-simulation/tests/test_examples.py +++ b/examples/nest-simulation/tests/test_examples.py @@ -62,7 +62,7 @@ def test_json_example(self): self.scaffold.compile() self._test_scaffold_results() results = self.scaffold.run_simulation("basal_activity") - self._test_simulation_results(results.spiketrains) + self._test_simulation_results(results.block.segments[0].spiketrains) def test_yaml_example(self): self.cfg = parse_configuration_file( @@ -72,7 +72,7 @@ def test_yaml_example(self): self.scaffold.compile() self._test_scaffold_results() results = self.scaffold.run_simulation("basal_activity") - self._test_simulation_results(results.spiketrains) + self._test_simulation_results(results.block.segments[0].spiketrains) def test_python_example(self): import scripts.guide_nest # noqa: F401 diff --git a/examples/neuron-simulation/tests/test_examples.py b/examples/neuron-simulation/tests/test_examples.py index 8ecd2208..0cefa8d1 100644 --- a/examples/neuron-simulation/tests/test_examples.py +++ b/examples/neuron-simulation/tests/test_examples.py @@ -67,7 +67,7 @@ def test_json_example(self): self.scaffold.compile() self._test_scaffold_results() results = self.scaffold.run_simulation("neuronsim") - self._test_simulation_results(results.analogsignals) + self._test_simulation_results(results.block.segments[0].analogsignals) def test_yaml_example(self): self.cfg = parse_configuration_file( @@ -77,7 +77,7 @@ def test_yaml_example(self): self.scaffold.compile() self._test_scaffold_results() results = self.scaffold.run_simulation("neuronsim") - self._test_simulation_results(results.analogsignals) + self._test_simulation_results(results.block.segments[0].analogsignals) def test_python_example(self): import scripts.guide_neuron # noqa: F401 diff --git a/packages/bsb-arbor/bsb_arbor/adapter.py b/packages/bsb-arbor/bsb_arbor/adapter.py index 72d874ef..b64021cb 100644 --- a/packages/bsb-arbor/bsb_arbor/adapter.py +++ b/packages/bsb-arbor/bsb_arbor/adapter.py @@ -25,11 +25,11 @@ class ArborSimulationData(SimulationData): Container class for simulation data. """ - def __init__(self, simulation): + def __init__(self, simulation, filename): """ Container class for simulation data. """ - super().__init__(simulation) + super().__init__(simulation, filename=filename) self.arbor_sim: arbor.simulation = None @@ -375,11 +375,13 @@ def __init__(self, comm=None): super().__init__(comm) self.simdata: dict[ArborSimulation, ArborSimulationData] = {} - def prepare(self, simulation: "ArborSimulation") -> ArborSimulationData: + def prepare( + self, simulation: "ArborSimulation", filename=None + ) -> ArborSimulationData: """ Prepares the arbor simulation engine with the given simulation. """ - simdata = self._create_simdata(simulation) + simdata = self._create_simdata(simulation, filename) try: context = arbor.context(arbor.proc_allocation(threads=simulation.threads)) if self.comm.get_size() > 1: @@ -466,8 +468,8 @@ def get_recipe(self, simulation, simdata=None): self._cache_devices(simulation, simdata) return ArborRecipe(simulation, simdata) - def _create_simdata(self, simulation): - self.simdata[simulation] = simdata = ArborSimulationData(simulation) + def _create_simdata(self, simulation, filename): + self.simdata[simulation] = simdata = ArborSimulationData(simulation, filename) self._assign_chunks(simulation, simdata) return simdata diff --git a/packages/bsb-core/bsb/cli/commands/_commands.py b/packages/bsb-core/bsb/cli/commands/_commands.py index cc9c402a..c3a6779f 100644 --- a/packages/bsb-core/bsb/cli/commands/_commands.py +++ b/packages/bsb-core/bsb/cli/commands/_commands.py @@ -229,13 +229,13 @@ def handler(self, context): level=0, ) try: - result = network.run_simulation(sim_name) + network.run_simulation(sim_name, output_filename=root / f"{uuid4()}.nio") except NodeNotFoundError as e: append = ", " if len(network.simulations) else "" append += ", ".join(f"'{name}'" for name in extra_simulations) errr.wrap(type(e), e, append=append) - else: - result.write(root / f"{uuid4()}.nio", "ow") + # else: + # result.write(root / f"{uuid4()}.nio", "ow") def get_options(self): return { diff --git a/packages/bsb-core/bsb/core.py b/packages/bsb-core/bsb/core.py index 9a1125f8..caf8ec25 100644 --- a/packages/bsb-core/bsb/core.py +++ b/packages/bsb-core/bsb/core.py @@ -449,7 +449,7 @@ def run_pipelines(self, fail_fast=True, pipelines=None): pool.schedule(pipelines) pool.execute() - def run_simulation(self, simulation_name: str): + def run_simulation(self, simulation_name: str, output_filename: str = None): """ Run a simulation starting from the default single-instance adapter. @@ -460,7 +460,7 @@ def run_simulation(self, simulation_name: str): adapter = get_simulation_adapter( simulation.simulator, comm=self._comm.get_communicator() ) - return adapter.simulate(simulation)[0] + return adapter.simulate(simulation, filename=output_filename)[0] def get_simulation(self, sim_name: str) -> Simulation: """ diff --git a/packages/bsb-core/bsb/simulation/adapter.py b/packages/bsb-core/bsb/simulation/adapter.py index b31e28a6..b365c303 100644 --- a/packages/bsb-core/bsb/simulation/adapter.py +++ b/packages/bsb-core/bsb/simulation/adapter.py @@ -66,7 +66,7 @@ def use_bar(self): class SimulationData: - def __init__(self, simulation: "Simulation", result=None): + def __init__(self, simulation: "Simulation", result=None, filename=None): self.chunks = None self.populations = dict() self.placement: dict[CellModel, PlacementSet] = { @@ -75,7 +75,7 @@ def __init__(self, simulation: "Simulation", result=None): self.connections = dict() self.devices = dict() if result is None: - result = SimulationResult(simulation) + result = SimulationResult(simulation, filename=filename) self.result: SimulationResult = result @@ -92,7 +92,7 @@ def __init__(self, comm=None): self._duration = None self.current_checkpoint = 0 - def simulate(self, *simulations, post_prepare=None): + def simulate(self, *simulations, post_prepare=None, filename=None): """ Simulate the given simulations. @@ -113,7 +113,7 @@ def simulate(self, *simulations, post_prepare=None): self._controllers.append(listener) for simulation in simulations: - data = self.prepare(simulation) + data = self.prepare(simulation, filename) alldata.append(data) for hook in simulation.post_prepare: hook(self, simulation, data) @@ -123,7 +123,7 @@ def simulate(self, *simulations, post_prepare=None): return self.collect(results) @abc.abstractmethod - def prepare(self, simulation): # pragma: nocover + def prepare(self, simulation, filename=None): # pragma: nocover """ Reset the simulation backend and prepare for the given simulation. diff --git a/packages/bsb-core/bsb/simulation/results.py b/packages/bsb-core/bsb/simulation/results.py index 1547d22e..fbf68025 100644 --- a/packages/bsb-core/bsb/simulation/results.py +++ b/packages/bsb-core/bsb/simulation/results.py @@ -9,22 +9,37 @@ class SimulationResult: - def __init__(self, simulation): - from neo import Block + def __init__(self, simulation, filename=None): + from neo import Block, io tree = simulation.__tree__() with contextlib.suppress(KeyError): del tree["post_prepare"] - self.block = Block(name=simulation.name, config=tree) + if filename: + self.filename = filename + self.name = simulation.name + io = io.NixIO(filename, mode="rw") + io.write(Block(name=self.name, nix_name=self.name, config=tree)) + for i, nixblock in enumerate(io.nix_file.blocks): + if self.name == nixblock.name: + self.block_id = i + io.close() + else: + self.block = Block( + name=simulation.name, nix_name=simulation.name, config=tree + ) + self.recorders = [] @property - def spiketrains(self): - return self.block.segments[0].spiketrains + def analogsignals(self): + if hasattr(self, "block"): + return self.block.segments[0].analogsignals @property - def analogsignals(self): - return self.block.segments[0].analogsignals + def spiketrains(self): + if hasattr(self, "block"): + return self.block.segments[0].spiketrains def add(self, recorder): self.recorders.append(recorder) @@ -36,21 +51,28 @@ def create_recorder(self, flush: typing.Callable[["neo.core.Segment"], None]): return recorder def flush(self): - from neo import Segment + from neo import Segment, io segment = Segment() - self.block.segments.append(segment) for recorder in self.recorders: try: recorder.flush(segment) except Exception: traceback.print_exc() warn("Recorder errored out!") + if hasattr(self, "filename"): + io = io.NixIO(self.filename, mode="rw") + block = io.nix_file.blocks[self.block_id] + io._write_segment(segment, block) + io.close() + else: + self.block.segments.append(segment) def write(self, filename, mode): from neo import io - io.NixIO(filename, mode=mode).write(self.block) + if hasattr(self, "block"): + io.NixIO(filename, mode=mode).write(self.block) class SimulationRecorder: diff --git a/packages/bsb-nest/bsb_nest/adapter.py b/packages/bsb-nest/bsb_nest/adapter.py index 7accd0ca..13109a13 100644 --- a/packages/bsb-nest/bsb_nest/adapter.py +++ b/packages/bsb-nest/bsb_nest/adapter.py @@ -51,14 +51,16 @@ def __init__(self, comm=None): self.loaded_modules = set() self._prev_chkpoint = 0 - def simulate(self, *simulations, post_prepare=None): + def simulate(self, *simulations, post_prepare=None, filename=None): try: self.reset_kernel() - return super().simulate(*simulations, post_prepare=post_prepare) + return super().simulate( + *simulations, post_prepare=post_prepare, filename=filename + ) finally: self.reset_kernel() - def prepare(self, simulation): + def prepare(self, simulation, filename=None): """ Prepare the simulation environment in NEST. @@ -80,7 +82,7 @@ def prepare(self, simulation): :rtype: bsb.simulation.adapter.SimulationData """ self.simdata[simulation] = SimulationData( - simulation, result=NestResult(simulation) + simulation, result=NestResult(simulation, filename) ) try: report("Installing NEST modules...", level=2) diff --git a/packages/bsb-neuron/bsb_neuron/adapter.py b/packages/bsb-neuron/bsb_neuron/adapter.py index d00ec34f..8402ca50 100644 --- a/packages/bsb-neuron/bsb_neuron/adapter.py +++ b/packages/bsb-neuron/bsb_neuron/adapter.py @@ -70,7 +70,7 @@ def engine(self): return engine - def prepare(self, simulation): + def prepare(self, simulation, filename=None): """ Prepare the simulation environment and data structures for running a NEURON simulation. @@ -90,7 +90,8 @@ def prepare(self, simulation): """ self.simdata[simulation] = NeuronSimulationData( - simulation, result=NeuronResult(simulation) + simulation, + result=NeuronResult(simulation, filename=filename), ) try: report("Preparing simulation", level=2)