From 032f977e5f3dc6c81533688432eff746b45a0d54 Mon Sep 17 00:00:00 2001
From: Filippo <52816133+filimarc@users.noreply.github.com>
Date: Wed, 13 Nov 2024 09:30:47 +0100
Subject: [PATCH] Added metadata to have more info when loading result nio
 files (#22)

Co-authored-by: danilobenozzo <danilo.benozzo@gmail.com>
---
 bsb_neuron/devices/current_clamp.py    | 12 +++++++++---
 bsb_neuron/devices/synapse_recorder.py | 13 ++++++++++---
 bsb_neuron/devices/voltage_clamp.py    | 12 +++++++++---
 bsb_neuron/devices/voltage_recorder.py | 12 +++++++++---
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/bsb_neuron/devices/current_clamp.py b/bsb_neuron/devices/current_clamp.py
index 16ee2b7..dcb4b0f 100644
--- a/bsb_neuron/devices/current_clamp.py
+++ b/bsb_neuron/devices/current_clamp.py
@@ -19,12 +19,18 @@ def implement(self, adapter, simulation, simdata):
                 for location in self.locations.get_locations(target):
                     if clamped:
                         warn(f"Multiple current clamps placed on {target}")
-                    self._add_clamp(simdata, location)
+                    self._add_clamp(
+                        simdata,
+                        location,
+                        name=self.name,
+                        cell_type=target.cell_model.name,
+                        cell_id=target.id,
+                    )
                     clamped = True
 
-    def _add_clamp(self, simdata, location):
+    def _add_clamp(self, simdata, location, **annotations):
         sx = location.arc(0.5)
         clamp = location.section.iclamp(
             x=sx, delay=self.before, duration=self.duration, amplitude=self.amplitude
         )
-        simdata.result.record(clamp._ref_i)
+        simdata.result.record(clamp._ref_i, **annotations)
diff --git a/bsb_neuron/devices/synapse_recorder.py b/bsb_neuron/devices/synapse_recorder.py
index d3627be..e3ed6ea 100644
--- a/bsb_neuron/devices/synapse_recorder.py
+++ b/bsb_neuron/devices/synapse_recorder.py
@@ -19,8 +19,15 @@ def implement(self, adapter, simulation, simdata):
                             not self.synapse_types
                             or synapse.synapse_name in self.synapse_types
                         ):
-                            _record_synaptic_current(simdata.result, synapse)
+                            _record_synaptic_current(
+                                simdata.result,
+                                synapse,
+                                name=self.name,
+                                cell_type=target.cell_model.name,
+                                cell_id=target.id,
+                                synapse_type=synapse.synapse_name,
+                            )
 
 
-def _record_synaptic_current(result, synapse):
-    result.record(synapse._pp._ref_i)
+def _record_synaptic_current(result, synapse, **annotations):
+    result.record(synapse._pp._ref_i, **annotations)
diff --git a/bsb_neuron/devices/voltage_clamp.py b/bsb_neuron/devices/voltage_clamp.py
index 3dc68f7..e5070f0 100644
--- a/bsb_neuron/devices/voltage_clamp.py
+++ b/bsb_neuron/devices/voltage_clamp.py
@@ -22,10 +22,16 @@ def implement(self, adapter, simulation, simdata):
             for location in self.locations.get_locations(target):
                 if clamped:
                     warnings.warn(f"Multiple voltage clamps placed on {target}")
-                self._add_clamp(simdata.results, location)
+                self._add_clamp(
+                    simdata.results,
+                    location,
+                    name=self.name,
+                    cell_type=target.cell_model.name,
+                    cell_id=target.id,
+                )
                 clamped = True
 
-    def _add_clamp(self, results, location):
+    def _add_clamp(self, results, location, **annotations):
         sx = location.arc(0.5)
         clamp = location.section.vclamp(
             voltage=self.voltage,
@@ -36,4 +42,4 @@ def _add_clamp(self, results, location):
                 if (v := getattr(self, k)) is not None
             },
         )
-        results.record(clamp)
+        results.record(clamp, **annotations)
diff --git a/bsb_neuron/devices/voltage_recorder.py b/bsb_neuron/devices/voltage_recorder.py
index 00c00b1..d04bfba 100644
--- a/bsb_neuron/devices/voltage_recorder.py
+++ b/bsb_neuron/devices/voltage_recorder.py
@@ -13,9 +13,15 @@ def implement(self, adapter, simulation, simdata):
         ).items():
             for target in pop:
                 for location in self.locations.get_locations(target):
-                    self._add_voltage_recorder(simdata.result, location)
+                    self._add_voltage_recorder(
+                        simdata.result,
+                        location,
+                        name=self.name,
+                        cell_type=target.cell_model.name,
+                        cell_id=target.id,
+                    )
 
-    def _add_voltage_recorder(self, results, location):
+    def _add_voltage_recorder(self, results, location, **annotations):
         section = location.section
         x = location.arc(0)
-        results.record(section(x)._ref_v)
+        results.record(section(x)._ref_v, **annotations)