` folder from the job " \
+ "artifacts using tvb-ext-unicore extension and update the paths to `simulation_data.npz` and " \
+ "`model_params.npz` accordingly.\n" \
+ "\n"
+
+ self.add_markdown_cell(intro)
+ plot_funct = self.vbi_plot_funct()
+ self.add_code_cell(plot_funct)
+ code = self.plot_timeseries()
+ self.add_code_cell(code)
+ self.add_code_cell(self.plot_from_batched_simulations())
+ self.add_code_cell(self.plot_single_simulation())
+
+ return self.notebook
+
+ @staticmethod
+ def vbi_plot_funct():
+ code = "from scipy import signal\n" \
+ "\n" \
+ "# The plotting helper is included inline for now, it will be replaced in the future.\n" \
+ "def plot_ts_pxx_jr(data, par, ax, method='welch', **kwargs):\n" \
+ " tspan = data['t']\n" \
+ " y = data['x']\n" \
+ " ax[0].plot(tspan, y.T, label='y1 - y2', **kwargs)\n" \
+ "\n" \
+ " if method == 'welch':\n" \
+ " freq, pxx = signal.welch(y, 1000/par['dt'], nperseg=y.shape[1]//2)\n" \
+ " else:\n" \
+ " freq, pxx = fft_signal(y, tspan / 1000)\n" \
+ " ax[1].plot(freq, pxx.T, **kwargs)\n" \
+ " ax[1].set_xlim(0, 50)\n" \
+ " ax[1].set_xlabel('frequency [Hz]')\n" \
+ " ax[0].set_xlabel('time [ms]')\n" \
+ " ax[0].set_ylabel('y1-y2')\n" \
+ " ax[0].margins(x=0)\n" \
+ " plt.tight_layout()\n"
+
+ return code
+
+ def plot_timeseries(self):
+ code = "import matplotlib.pyplot as plt\n" \
+ "import numpy as np\n" \
+ "from xai_components.serialization import *\n" \
+ "\n" \
+ "data = np.load('{data_path}')\n" \
+ "params = load_params_npz('{params_path}')\n"
+
+ inputs = self._prepare_component_inputs()
+ return code.format(**inputs)
+
+ @staticmethod
+ def plot_from_batched_simulations():
+ code = "# Run this cell for CuPy backend: the saved file contains all simulations in one batch.\n" \
+ "ts0 = data['x'][:, :, 0].T\n" \
+ "data0 = {'t': data['t'], 'x': ts0}\n" \
+ "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n" \
+ "plot_ts_pxx_jr(data0, params, ax, alpha=0.6, lw=1)\n" \
+ "plt.tight_layout()\n"
+ return code
+
+ @staticmethod
+ def plot_single_simulation():
+ code = "# Run this cell for C++/Numba backend: the saved file contains one simulation only.\n" \
+ "fig, ax = plt.subplots(1, 2, figsize=(10, 3))\n" \
+ "plot_ts_pxx_jr(data, params, ax, alpha=0.6, lw=1)\n" \
+ "plt.tight_layout()\n"
+ return code
+
+ def _prepare_component_inputs(self):
+ base_root = os.path.join(get_base_dir_web(), "output")
+ output_dir = os.path.join(base_root, "output" + f"_{self.xircuits_filename}")
+
+ params_path = os.path.join(output_dir, "model_params.npz").replace("\\", "/") if IS_WINDOWS \
+ else os.path.join(output_dir, "model_params.npz")
+
+ data_path = os.path.join(output_dir, "simulation_data.npz").replace("\\", "/") if IS_WINDOWS \
+ else os.path.join(output_dir, "simulation_data.npz")
+ return {
+ "params_path": params_path,
+ "data_path": data_path,
+ }
+
+ def edit_cell(self):
+ return True
+
def determine_component_class(component_name, component_path):
component_module = importlib.import_module(component_path.replace('/', '.')[:-3])
component_class = getattr(component_module, component_name)
diff --git a/tvbextxircuits/utils.py b/tvbextxircuits/utils.py
index a6d97ec4..7a64e094 100644
--- a/tvbextxircuits/utils.py
+++ b/tvbextxircuits/utils.py
@@ -1,4 +1,4 @@
-import json
+import json5
from pathlib import Path
from jupyter_core.paths import jupyter_config_dir
@@ -55,7 +55,7 @@ def get_user_settings():
settings_path = os.path.join(data_dir, 'lab', 'user-settings', 'tvb-ext-xircuits', 'settings.jupyterlab-settings')
if os.path.exists(settings_path):
with open(settings_path, 'r', encoding='utf-8') as f:
- settings = json.load(f)
+ settings = json5.load(f)
else:
settings = {}
diff --git a/xai_components/base_tvb.py b/xai_components/base_tvb.py
index 36a7f9ad..dff161c9 100644
--- a/xai_components/base_tvb.py
+++ b/xai_components/base_tvb.py
@@ -26,3 +26,9 @@ class ComponentWithWidget(TVBComponent):
"""
Used to flag a component that has an associate widget to be displayed in Xircuits UI for interactive setup.
"""
+
+class ComponentWithViewer(Component):
+ """"
+ Marker for components that expose an 'Open Viewer' functionality, but are not from TVB world.
+ """
+ pass
diff --git a/xai_components/serialization.py b/xai_components/serialization.py
new file mode 100644
index 00000000..c28cb733
--- /dev/null
+++ b/xai_components/serialization.py
@@ -0,0 +1,58 @@
+import json
+import numpy as np
+
+DATATYPES_KEY = "__datatypes__"
+
+def save_params_npz(params, path):
+ """
+ Save a heterogeneous parameter dict to a single npz file.
+ """
+ payload = {}
+ datatypes = {}
+
+ for key, val in params.items():
+ if isinstance(val, np.ndarray):
+ payload[key] = val
+ datatypes[key] = "ndarray"
+ elif isinstance(val, (bool, int, float)):
+ payload[key] = np.array(val)
+ datatypes[key] = "scalar"
+ elif isinstance(val, str):
+ payload[key] = np.array(val)
+ datatypes[key] = "str"
+ elif val is None:
+ payload[key] = np.array(0)
+ datatypes[key] = "none"
+ else:
+ raise TypeError(f"Unsupported param type for '{key}': {type(val)}")
+ payload[DATATYPES_KEY] = np.array(json.dumps(datatypes), dtype=np.str_)
+ np.savez(path, **payload)
+
+def load_params_npz(path):
+ """
+ Load parameters saved by save_params_npz().
+ """
+ with np.load(path, allow_pickle=False) as npzfile:
+ if DATATYPES_KEY in npzfile.files:
+ datatype_arr = json.loads(str(npzfile[DATATYPES_KEY].item()))
+ else:
+ datatype_arr = {}
+
+ output= {}
+ for key in npzfile.files:
+ if key == DATATYPES_KEY:
+ continue
+
+ arr = npzfile[key]
+ datatype = datatype_arr.get(key)
+
+ if datatype == "ndarray":
+ output[key] = arr
+ elif datatype in ("scalar", "str"):
+ output[key] = arr.item()
+ elif datatype == "none":
+ output[key] = None
+ else:
+ output[key] = arr
+
+ return output
\ No newline at end of file
diff --git a/xai_components/settings.py b/xai_components/settings.py
new file mode 100644
index 00000000..d39815b7
--- /dev/null
+++ b/xai_components/settings.py
@@ -0,0 +1,23 @@
+import os
+from joblib import Memory
+
+# Paths
+CACHE_DIR = os.path.join(os.path.expanduser('~'), '.cache-xircuits')
+
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# Cache
+def _get_vbi_version():
+ try:
+ import vbi
+ return getattr(vbi, "__version__", "unknown")
+ except Exception:
+ return "unknown"
+
+def create_memory():
+ vbi_version = _get_vbi_version()
+ cache_dir = os.path.join(CACHE_DIR, f"vbi-{vbi_version}")
+ os.makedirs(cache_dir, exist_ok=True)
+ return Memory(cache_dir, verbose=2)
+
+memory = create_memory()
\ No newline at end of file
diff --git a/xai_components/xai_vbi_config_inference/__init__.py b/xai_components/xai_vbi_config_inference/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/xai_components/xai_vbi_config_inference/arguments/configinference.json b/xai_components/xai_vbi_config_inference/arguments/configinference.json
new file mode 100644
index 00000000..183ade8b
--- /dev/null
+++ b/xai_components/xai_vbi_config_inference/arguments/configinference.json
@@ -0,0 +1,12 @@
+{
+ "description": "\n
Configuration Inference
\n
Builds the configuration inputs required to run an inference workflow.
\n
It defines the priors, samples them to obtain the parameter distribution (theta) and assembles a feature dictionary used for later feature-extraction.
\n
\n
Generic Hopfield model
\n
The Hopfield neural network is a discrete time dynamical system composed\n of multiple binary nodes, with a connectivity matrix built from a\n predetermined set of patterns. The update, inspired from the spin-glass\n model (used to describe magnetic properties of dilute alloys), is based on\n a random scanning of every node. The existence of a fixed point dynamics\n is guaranteed by a Lyapunov function. The Hopfield network is expected to\n have those multiple patterns as attractors (multistable dynamical system).\n When the initial conditions are close to one of the 'learned' patterns,\n the dynamical system is expected to relax on the corresponding attractor.\n A possible output of the system is the final attractive state (interpreted\n as an associative memory).\n\n Various extensions of the initial model have been proposed, among which a\n noiseless and continuous version [Hopfield 1984] having a slightly\n different Lyapunov function, but essentially the same dynamical\n properties, with more straightforward physiological interpretation. A\n continuous Hopfield neural network (with a sigmoid transfer function) can\n indeed be interpreted as a network of neural masses with every node\n corresponding to the mean field activity of a local brain region, with\n many bridges with the Wilson Cowan model [WC_1972].
\n
\n\n\n| [Hopfield1982] | Hopfield, J. J. *Neural networks and physical systems with emergent collective computational abilities*, Proc. Nat. Acad. Sci. (USA) 79, 2554-2558, 1982. |
\n\n
\n\n
\n\n\n| [Hopfield1984] | Hopfield, J. J. *Neurons with graded response have collective computational properties like those of two-sate neurons*, Proc. Nat. Acad. Sci. (USA) 81, 3088-3092, 1984. |
\n\n
'heun' or 'euler'.\n",
+ "output": "Output directory.\n",
+ "initial_state" : "Initial state of the system of shape (`num_nodes`, `num_sim`).\n",
+ "same_initial_state" : "If True, same initial condition for all simulations\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/arguments/jrsdecpp.json b/xai_components/xai_vbi_models/arguments/jrsdecpp.json
new file mode 100644
index 00000000..6d8fc08a
--- /dev/null
+++ b/xai_components/xai_vbi_models/arguments/jrsdecpp.json
@@ -0,0 +1,29 @@
+{
+ "description": "\n
Jansen-Rit model
\n
The Jansen-Rit neural mass model [Jansen1995] has been widely used to simulate physiological signals from various recording methods like intracranial LFPs, and scalp MEG/EEG recordings.
\n
For example, it has been shown to recreate responses similar to evoked-related potentials after a series of impulse stimulations [David2003], [David_etal06], generating high-alpha and low-beta oscillations (with added recurrent inhibitory connections and spike-rate modulation) [Moran2007], and also seizure patterns similar to those seen in temporal lobe epilepsy [Wendling2001]. This biologically motivated model comprises of three main populations of neurons: excitatory pyramidal neurons, inhibitory interneurons, and excitatory interneurons. These populations interact with each other through synaptic connections, forming a feedback loop that produces oscillatory activity governed by a set of nonlinear ordinary differential equations [JansenRit], [David2003], [Kazemi2022].\n
\n\n\n| [Jansen1995] | Jansen, B. H., & Rit, V. G. (1995). Electroencephalogram and visual evoked potential generation in a mathematical model of coupled cortical columns. Biological Cybernetics, 73(4), 357-366. |
\n\n
\n\n
\n\n\n| [David2003] | David, O., & Friston, K. J. (2003). A neural mass model for MEG/EEG: coupling and neuronal dynamics. NeuroImage, 20(3), 1743-1755. https://doi.org/10.1016/j.neuroimage.2003.07.015 |
\n\n
\n\n
\n\n\n| [David_etal06] | David, O., Kiebel, S. J., Harrison, L. M., Mattout, J., Kilner, J. M., & Friston, K. J. (2006). Dynamic causal modeling of evoked responses in EEG and MEG. NeuroImage, 30(4), 1255-1272. https://doi.org/10.1016/j.neuroimage.2005.10.045. |
\n\n
\n\n
\n\n\n| [Moran2007] | Moran, R. J., Kiebel, S. J., Stephan, K. E., Reilly, R. B., Daunizeau, J., & Friston, K. J. (2007). A neural mass model of spectral responses in electrophysiology. NeuroImage, 37(3), 706-720. https://doi.org/10.1016/j.neuroimage.2007.05.032. |
\n\n
\n\n
\n\n\n| [Wendling2001] | Wendling, F., Bartolomei, F., Bellanger, J.-J., & Chauvel, P. (2001). Interpretation of interdependencies in epileptic signals using a macroscopic physiological model of the EEG. Clinical Neurophysiology, 112(7), 1201-1218. |
\n\n
\n\n
\n\n\n| [JansenRit] | Jansen, B. H., & Rit, V. G. (1995). Electroencephalogram and visual evoked potential generation in a mathematical model of coupled cortical columns. Biological Cybernetics, 73, 357-366. |
\n\n
\n\n
\n\n\n| [Kazemi2022] | Kazemi, S., & Jamali, Y. (2022). On the influence of input triggering on the dynamics of the Jansen-Rit oscillators network. arXiv preprint arXiv:2202.06634. |
\n\n
\n
'heun' or 'euler'.\n",
+ "t_end" : "End time of the simulation.\n",
+ "t_transition" : "Transient period to discard from the beginning of the simulation.\n",
+ "output": "Output directory.\n",
+ "RECORD_AVG": "If True, stores large time series during the simulations.\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/arguments/jrsdecupy.json b/xai_components/xai_vbi_models/arguments/jrsdecupy.json
new file mode 100644
index 00000000..31db893d
--- /dev/null
+++ b/xai_components/xai_vbi_models/arguments/jrsdecupy.json
@@ -0,0 +1,33 @@
+{
+ "description": "\n
Jansen-Rit model
\n
The Jansen-Rit neural mass model [Jansen1995] has been widely used to simulate physiological signals from various recording methods like intracranial LFPs, and scalp MEG/EEG recordings.
\n
For example, it has been shown to recreate responses similar to evoked-related potentials after a series of impulse stimulations [David2003], [David_etal06], generating high-alpha and low-beta oscillations (with added recurrent inhibitory connections and spike-rate modulation) [Moran2007], and also seizure patterns similar to those seen in temporal lobe epilepsy [Wendling2001]. This biologically motivated model comprises of three main populations of neurons: excitatory pyramidal neurons, inhibitory interneurons, and excitatory interneurons. These populations interact with each other through synaptic connections, forming a feedback loop that produces oscillatory activity governed by a set of nonlinear ordinary differential equations [JansenRit], [David2003], [Kazemi2022].\n
\n\n\n| [Jansen1995] | Jansen, B. H., & Rit, V. G. (1995). Electroencephalogram and visual evoked potential generation in a mathematical model of coupled cortical columns. Biological Cybernetics, 73(4), 357-366. |
\n\n
\n\n
\n\n\n| [David2003] | David, O., & Friston, K. J. (2003). A neural mass model for MEG/EEG: coupling and neuronal dynamics. NeuroImage, 20(3), 1743-1755. https://doi.org/10.1016/j.neuroimage.2003.07.015 |
\n\n
\n\n
\n\n\n| [David_etal06] | David, O., Kiebel, S. J., Harrison, L. M., Mattout, J., Kilner, J. M., & Friston, K. J. (2006). Dynamic causal modeling of evoked responses in EEG and MEG. NeuroImage, 30(4), 1255-1272. https://doi.org/10.1016/j.neuroimage.2005.10.045. |
\n\n
\n\n
\n\n\n| [Moran2007] | Moran, R. J., Kiebel, S. J., Stephan, K. E., Reilly, R. B., Daunizeau, J., & Friston, K. J. (2007). A neural mass model of spectral responses in electrophysiology. NeuroImage, 37(3), 706-720. https://doi.org/10.1016/j.neuroimage.2007.05.032. |
\n\n
\n\n
\n\n\n| [Wendling2001] | Wendling, F., Bartolomei, F., Bellanger, J.-J., & Chauvel, P. (2001). Interpretation of interdependencies in epileptic signals using a macroscopic physiological model of the EEG. Clinical Neurophysiology, 112(7), 1201-1218. |
\n\n
\n\n
\n\n\n| [JansenRit] | Jansen, B. H., & Rit, V. G. (1995). Electroencephalogram and visual evoked potential generation in a mathematical model of coupled cortical columns. Biological Cybernetics, 73, 357-366. |
\n\n
\n\n
\n\n\n| [Kazemi2022] | Kazemi, S., & Jamali, Y. (2022). On the influence of input triggering on the dynamics of the Jansen-Rit oscillators network. arXiv preprint arXiv:2202.06634. |
\n\n
\n
\n
Montbrio-Pazo-Roxin model
\n
The exact macroscopic dynamics of a specific brain region (represented as a node in the network) can be analytically derived in the thermodynamic limit of infinitely all-to-all coupled spiking neurons [Montbrio2015] or `theta` neuron representation [Byrne2020next]. By assuming a Lorentzian distribution on excitabilities in large ensembles of quadratic integrate-and-fire neurons with synaptic weights J and a half-width `delta` centered at `eta`, the macroscopic dynamics has been derived in terms of the collective firing activity and mean membrane potential [Montbrio2015].
\n
\n\n\n| [Montbrio2015] | Montbrio, E., et al. (2015). Macroscopic description for networks of spiking neurons. Physical Review X, 5(2), 021028. |
\n\n
\n\n
\n\n\n| [Byrne2020next] | Byrne, A., et al. (2020). Next generation neural mass models. Journal of Neuroscience Methods, 340, 108746. |
\n\n
r and v.\n",
+ "RECORD_BOLD": "If True, stores the simulated BOLD signal.\n",
+ "RECORD_AVG_r": "If True, stores the node-averaged firing rate.\n",
+ "num_sim": "Number of simulations to run (batch size).\n",
+ "method" : "Numerical integration method. Supported values: 'heun' or 'euler'.\n",
+ "engine" : "Supported values: `cpu` or `gpu`\n",
+ "seed" : "Random seed for reproducibility.\n",
+ "dtype" : "Data type to use for the simulation, `float` for `float64` or `f` for `float32`.\n",
+ "initial_state" : "Initial state of the system of shape (`num_nodes`, `num_sim`).\n",
+ "same_initial_state" : "If True, use the same initial condition for all simulations.\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/arguments/vepsdecpp.json b/xai_components/xai_vbi_models/arguments/vepsdecpp.json
new file mode 100644
index 00000000..995e7b35
--- /dev/null
+++ b/xai_components/xai_vbi_models/arguments/vepsdecpp.json
@@ -0,0 +1,20 @@
+{
+ "description": "\n
Virtual Epileptic Patient model
\n
In personalized whole-brain network modeling of epilepsy spread [Jirsa2017], the dynamics of each brain region are governed by the Epileptor model [Jirsa2014]. The Epileptor model provides a comprehensive description of epileptic seizures, encompassing the complete taxonomy of system bifurcations to simultaneously reproduce the dynamics of seizure onset, progression, and termination [Saggio2020]. The full Epileptor model comprises five state variables that couple two oscillatory dynamical systems operating on three different time scales [Jirsa2014]. Then motivated by Synergetic theory [Haken1997], [JirsaHaken1997] and under time-scale separation [Proix2014], the fast variables rapidly collapse on the slow manifold [McIntoshJirsa2019], whose dynamics is governed by the slow variable.\n
\n
\n\n\n| [Jirsa2017] | Jirsa, V.K.; Proix, T.; Perdikis, D.; Woodman, M.M.; Wang, H.; Gonzalez-Martinez, J.; Bernard, C.; Bénar, C.; Guye, M.; Chauvel, P.; Bartolomei, F. (2017). The Virtual Epileptic Patient: Individualized whole-brain models of epilepsy spread. NeuroImage, 145, 377-388, doi:https://doi.org/10.1016/j.neuroimage.2016.04.049. |
\n\n
\n\n
\n\n\n| [Jirsa2014] | Jirsa, Viktor K.; Stacey, William C.; Quilichini, Pascale P.; Ivanov, Anton I.; Bernard, Christophe (2014). On the nature of seizure dynamics. Brain, 137(8), 2210-2230, doi:10.1093/brain/awu133. |
\n\n
\n\n
\n\n\n| [Saggio2020] | Saggio, Maria Luisa; Crisp, Dakota; Scott, Jared M; Karoly, Philippa; Kuhlmann, Levin; Nakatani, Mitsuyoshi; Murai, Tomohiko; Dümpelmann, Matthias; Schulze-Bonhage, Andreas; Ikeda, Akio; Cook, Mark; Gliske, Stephen V; Lin, Jack; Bernard, Christophe; Jirsa, Viktor; Stacey, William C (2020). A taxonomy of seizure dynamotypes. eLife, 9, e55632, doi:10.7554/eLife.55632. |
\n\n
\n\n
\n\n\n| [Haken1997] | \nHaken, Herman (1977). Synergetics. Physics Bulletin, 28(9), 412. |
\n\n
\n\n
\n\n\n| [JirsaHaken1997] | Jirsa, Viktor K; Haken, Hermann (1997). A derivation of a macroscopic field theory of the brain from the quasi-microscopic neural dynamics. Physica D: Nonlinear Phenomena, 99(4), 503-526. |
\n\n
\n\n
\n\n\n| [Proix2014] | Proix, Timothée; Bartolomei, Fabrice; Chauvel, Patrick; Bernard, Christophe; Jirsa, Viktor K. (2014). Permittivity Coupling across Brain Regions Determines Seizure Recruitment in Partial Epilepsy. Journal of Neuroscience, 34(45), 15009-15021, doi:10.1523/JNEUROSCI.1570-14.2014. |
\n\n
\n\n
\n\n\n| [McIntoshJirsa2019] | \nMcIntosh, Anthony R.; Jirsa, Viktor K. (2019). The hidden repertoire of brain dynamics and dysfunction. Network Neuroscience, 3(4), 994-1008, doi:10.1162/netn_a_00107. |
\n\n
\n
'heun' or 'euler'.\n",
+ "output": "Output directory.\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/ghb_sde_cupy.py b/xai_components/xai_vbi_models/ghb_sde_cupy.py
new file mode 100644
index 00000000..e66398d0
--- /dev/null
+++ b/xai_components/xai_vbi_models/ghb_sde_cupy.py
@@ -0,0 +1,39 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+from typing import Union
+
+@xai_component(color='rgb(101, 179, 46)')
+class GHBSdeCupy(Component):
+ G: InArg[Union[float, list]]
+ t_cut: InArg[float]
+ dt: InArg[float]
+ eta: InArg[int]
+ num_sim: InArg[int]
+ sigma: InArg[float]
+ seed: InArg[int]
+ decimate: InArg[[int]]
+ omega: InArg[Union[float, list]]
+ t_end: InArg[float]
+ engine: InArg[str]
+ weights: InArg[list]
+ dtype: InArg[str]
+ method: InArg[str]
+ output: InArg[str]
+ initial_state: InArg[list]
+ same_initial_state: InArg[bool]
+
+ model: OutArg[any]
+ time_series_key: OutArg[dict]
+
+ def execute(self, ctx):
+ from vbi.models.cupy.ghb import GHB_sde
+
+ keys = ["G", "t_cut", "dt", "eta", "num_sim", "sigma", "seed", "decimate", "omega", "t_end", "engine",
+ "weights", "dtype", "method", "output", "initial_state", "same_initial_state"]
+ params = {}
+ for key in keys:
+ param = getattr(self, key, None)
+ if param.value is not None:
+ params[key] = param.value
+
+ self.model.value = GHB_sde(par=params)
+ self.time_series_key.value = {"t": "t", "x": "bold"}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/jr_sde_cpp.py b/xai_components/xai_vbi_models/jr_sde_cpp.py
new file mode 100644
index 00000000..7c6beb92
--- /dev/null
+++ b/xai_components/xai_vbi_models/jr_sde_cpp.py
@@ -0,0 +1,47 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+from typing import Union
+
+@xai_component(color='rgb(101, 179, 46)')
+class JRSdeCpp(Component):
+ G: InArg[Union[float, list]]
+ A: InArg[Union[float, list]]
+ B: InArg[Union[float, list]]
+ a: InArg[Union[float, list]]
+ b: InArg[Union[float, list]]
+ noise_mu: InArg[Union[float, list]]
+ noise_std: InArg[Union[float, list]]
+ vmax: InArg[float]
+ v0: InArg[Union[float, list]]
+ r: InArg[Union[float, list]]
+ initial_state: InArg[list]
+ weights: InArg[list]
+ C0: InArg[Union[float, list]]
+ C1: InArg[Union[float, list]]
+ C2: InArg[Union[float, list]]
+ C3: InArg[Union[float, list]]
+ noise_seed: InArg[int] #?type
+ seed: InArg[int]
+ dt: InArg[float]
+ method: InArg[str]
+ t_transition: InArg[float]
+ t_end: InArg[float]
+ output: InArg[str]
+ RECORD_AVG: InArg[bool]
+
+ model: OutArg[any]
+ time_series_key: OutArg[dict]
+
+ def execute(self, ctx):
+ from vbi.models.cpp.jansen_rit import JR_sde
+
+ keys = ["noise_seed", "seed", "G", "weights", "A", "B", "a", "b", "noise_mu", "noise_std", "vmax", "v0", "r",
+ "C0", "C1", "C2", "C3", "dt", "method", "t_transition", "t_end", "output", "RECORD_AVG",
+ "initial_state"]
+ params = {}
+ for key in keys:
+ param = getattr(self, key, None)
+ if param.value is not None:
+ params[key] = param.value
+
+ self.model.value = JR_sde(par=params)
+ self.time_series_key.value = {"t": "t", "x": "x"}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/jr_sde_cupy.py b/xai_components/xai_vbi_models/jr_sde_cupy.py
new file mode 100644
index 00000000..e5a64289
--- /dev/null
+++ b/xai_components/xai_vbi_models/jr_sde_cupy.py
@@ -0,0 +1,52 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+from typing import Union
+
+@xai_component(color='rgb(101, 179, 46)')
+class JRSdeCupy(Component):
+ G: InArg[Union[float, list]]
+ A: InArg[Union[float, list]]
+ B: InArg[Union[float, list]]
+ v: InArg[Union[float, list]]
+ r: InArg[Union[float, list]]
+ v0: InArg[Union[float, list]]
+ vmax: InArg[float]
+ C0: InArg[Union[float, list]]
+ C1: InArg[Union[float, list]]
+ C2: InArg[Union[float, list]]
+ C3: InArg[Union[float, list]]
+ a: InArg[Union[float, list]]
+ b: InArg[Union[float, list]]
+ mu: InArg[Union[float, list]]
+ noise_amp: InArg[Union[float, list]]
+ decimate: InArg[[int]]
+ dt: InArg[float]
+ t_end: InArg[float]
+ t_cut: InArg[float]
+ engine: InArg[str]
+ method: InArg[str]
+ num_sim: InArg[int]
+ weights: InArg[list]
+ dtype: InArg[str]
+ seed: InArg[int]
+ initial_state: InArg[list]
+ same_initial_state: InArg[bool]
+ same_noise_per_sim: InArg[bool]
+
+ model: OutArg[any]
+ time_series_key: OutArg[dict]
+
+ def execute(self, ctx):
+ from vbi.models.cupy.jansen_rit import JR_sde
+
+ keys = ["G", "A", "B", "v", "r", "v0", "vmax", "C0", "C1", "C2", "C3", "a", "b", "mu", "noise_amp",
+ "decimate", "dt", "t_end", "t_cut", "engine", "method", "num_sim", "weights", "dtype", "seed",
+ "initial_state", "same_initial_state", "same_noise_per_sim"]
+ params = {}
+ for key in keys:
+ param = getattr(self, key, None)
+ if param.value is not None:
+ params[key] = param.value
+
+ self.model.value = JR_sde(par=params)
+ self.time_series_key.value = {"t": "t", "x": "x"}
+
diff --git a/xai_components/xai_vbi_models/mpr_sde_cupy.py b/xai_components/xai_vbi_models/mpr_sde_cupy.py
new file mode 100644
index 00000000..a02e0a65
--- /dev/null
+++ b/xai_components/xai_vbi_models/mpr_sde_cupy.py
@@ -0,0 +1,57 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+from typing import Union
+
+@xai_component(color='rgb(101, 179, 46)')
+class MPRSdeCupy(Component):
+ G: InArg[Union[float, list]]
+ dt: InArg[float]
+ dt_bold: InArg[float]
+ J: InArg[Union[float, list]]
+ eta: InArg[Union[float, list]]
+ tau: InArg[float]
+ delta: InArg[float]
+ tr: InArg[float]
+ noise_amp: InArg[Union[float, list]]
+ same_noise_per_sim: InArg[bool]
+ sti_apply: InArg[bool]
+ iapp: InArg[Union[float, list]]
+ t_start: InArg[float]
+ t_cut: InArg[float]
+ t_end: InArg[float]
+ num_nodes: InArg[int]
+ weights: InArg[list]
+ rv_decimate: InArg[int]
+ output: InArg[str]
+ RECORD_RV: InArg[bool]
+ RECORD_BOLD: InArg[bool]
+ RECORD_AVG_r: InArg[bool]
+ num_sim: InArg[int]
+ method: InArg[str]
+ engine: InArg[str]
+ seed: InArg[int]
+ dtype: InArg[str]
+ initial_state: InArg[list]
+ same_initial_state: InArg[bool]
+
+ model: OutArg[any]
+ time_series_key: OutArg[dict]
+
+ def __init__(self):
+ super().__init__()
+
+ def execute(self, ctx) -> None:
+ from vbi.models.cupy.mpr import MPR_sde
+
+ keys = ["G", "dt", "dt_bold", "J", "eta", "tau", "delta", "tr", "noise_amp", "same_noise_per_sim", "sti_apply",
+ "iapp", "t_start", "t_cut", "t_end", "num_nodes", "weights", "rv_decimate", "output", "RECORD_RV",
+ "RECORD_BOLD", "RECORD_AVG_r", "num_sim", "method", "engine", "seed", "dtype", "initial_state",
+ "same_initial_state"]
+ params = {}
+ for key in keys:
+ param = getattr(self, key, None)
+ if param.value is not None:
+ params[key] = param.value
+
+ self.model.value = MPR_sde(par=params)
+ self.time_series_key.value = {"t": "rv_t", "x": "rv_d"} if self.model.value.RECORD_RV \
+ else {"t": "fmri_t", "x": "fmri_d"}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_models/vep_sde_cpp.py b/xai_components/xai_vbi_models/vep_sde_cpp.py
new file mode 100644
index 00000000..1f391027
--- /dev/null
+++ b/xai_components/xai_vbi_models/vep_sde_cpp.py
@@ -0,0 +1,36 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+from typing import Union
+
+@xai_component(color='rgb(101, 179, 46)')
+class VEPSdeCpp(Component):
+ G: InArg[Union[float, list]]
+ seed: InArg[int]
+ initial_state: InArg[list]
+ weights: InArg[list]
+ tau: InArg[float]
+ eta: InArg[Union[int, float, list]]
+ noise_sigma: InArg[float]
+ iext: InArg[Union[int, float, list]]
+ dt: InArg[float]
+ tend: InArg[float]
+ tcut: InArg[float]
+ noise_seed: InArg[int]
+ record_step: InArg[int]
+ method: InArg[str]
+ output: InArg[str]
+
+ model: OutArg[any]
+ time_series_key: OutArg[dict]
+
+ def execute(self, ctx) -> None:
+ from vbi.models.cpp.vep import VEP_sde
+
+ keys = ["G", "seed", "initial_state", "weights", "tau", "eta", "noise_sigma", "iext", "dt", "tend", "tcut",
+ "noise_seed", "record_step", "method", "output"]
+ params = {}
+ for key in keys:
+ param = getattr(self, key, None)
+ if param.value is not None:
+ params[key] = param.value
+ self.model.value = VEP_sde(par=params)
+ self.time_series_key.value = {"t": "t", "x": "x"}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_sample_posterior/__init__.py b/xai_components/xai_vbi_sample_posterior/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/xai_components/xai_vbi_sample_posterior/arguments/sampleposterior.json b/xai_components/xai_vbi_sample_posterior/arguments/sampleposterior.json
new file mode 100644
index 00000000..e777d054
--- /dev/null
+++ b/xai_components/xai_vbi_sample_posterior/arguments/sampleposterior.json
@@ -0,0 +1,10 @@
+{
+ "description": "\n
Sample Posterior
\n
Samples parameter values from the trained posterior distribution, conditioned on a selected observed feature vector.
\n
sample() method).\n",
+ "X_scaled": "Observed data point of shape (1, n_features) or (n_features,). This is the target observation for which we want to infer parameters.\n",
+ "num_samples" : "Number of posterior samples to draw.\n",
+ "obs_idx" : "Row index used to select the observation from the standardized feature matrix. Default is 0.\n",
+ "output_dir" : "Output folder path where resulted samples are saved (used by viewers/plots).\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_sample_posterior/sample_posterior.py b/xai_components/xai_vbi_sample_posterior/sample_posterior.py
new file mode 100644
index 00000000..389687f7
--- /dev/null
+++ b/xai_components/xai_vbi_sample_posterior/sample_posterior.py
@@ -0,0 +1,43 @@
+from xai_components.base import xai_component, InArg, OutArg
+import torch
+from vbi.sbi_inference import Inference
+import os
+from xai_components.base_tvb import ComponentWithViewer
+from tvbextxircuits.logger.builder import get_logger
+
+LOGGER = get_logger(__name__)
+
+
+@xai_component(color="rgb(220, 5, 45)")
+class SamplePosterior(ComponentWithViewer):
+ posterior: InArg[any]
+ X_scaled: InArg[any]
+ num_samples: InArg[int]
+ obs_idx: InArg[int]
+ output_dir: InArg[str]
+
+ samples: OutArg[torch.Tensor]
+
+ def __init__(self):
+ super().__init__()
+ self.obs_idx.value = 0
+
+ def execute(self, ctx):
+
+ x_idx_st = self.X_scaled.value[self.obs_idx.value,:]
+
+ # Sample Posterior
+ obj = Inference()
+ samples = obj.sample_posterior(x_idx_st, self.num_samples.value, self.posterior.value)
+
+ self.persists_artifacts(self.output_dir.value, samples)
+
+ self.samples.value = samples
+
+ LOGGER.info(f"Sample Posterior: {self.samples.value}")
+
+ @staticmethod
+ def persists_artifacts(output_dir, samples):
+ # Store the resulted samples for plotting
+ path = os.path.join(output_dir, "samples.pt")
+ torch.save(samples, path)
\ No newline at end of file
diff --git a/xai_components/xai_vbi_simulation_runner/__init__.py b/xai_components/xai_vbi_simulation_runner/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/xai_components/xai_vbi_simulation_runner/arguments/simulationrunner.json b/xai_components/xai_vbi_simulation_runner/arguments/simulationrunner.json
new file mode 100644
index 00000000..5d0cc260
--- /dev/null
+++ b/xai_components/xai_vbi_simulation_runner/arguments/simulationrunner.json
@@ -0,0 +1,13 @@
+{
+ "description": "\n
Simulation Runner
\n
Runs the selected VBI model for a batch of parameter samples (theta) using the chosen backend (cpp, cupy or numba). The component updates model parameters from theta, simulates the requested output signal (selected via time_series_key), and extracts summary features defined in cfg. The output is a feature matrix used for posterior training.
\n
cpp, cupy, numba.\n",
+ "model": "The instantiated VBI model used for simulation (e.g., Jansen-Rit, MPR).\n",
+ "theta": "A tensor/array of parameter samples with shape (N, D), where N is the number of samples/simulations and D is the dimension of the parameter space\n",
+ "theta_names": "List of parameter names (length D). Must match the column order of theta.\n",
+ "cfg": "Feature-extraction configuration dictionary used to compute summary statistics from the simulated time series.\n",
+ "num_workers": "Number of parallel workers. Used for running many independent simulations concurrently (for backends that do not batch simulations, e.g., cpp, numba).\n",
+ "time_series_key": "Dictionary specifying which outputs to read from the model result, e.g. {'t': 'rv_t', 'x': 'rv_d'}. Use this to choose the observed signal (RV, BOLD/fMRI, etc.) returned by the model.\n",
+ "output_dir": "Output folder path where model parameters and simulation data are saved (used by viewers/plots).\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_simulation_runner/simulation_runner.py b/xai_components/xai_vbi_simulation_runner/simulation_runner.py
new file mode 100644
index 00000000..123fd32e
--- /dev/null
+++ b/xai_components/xai_vbi_simulation_runner/simulation_runner.py
@@ -0,0 +1,182 @@
+from xai_components.base import xai_component, InArg, OutArg
+import numpy as np
+import torch
+from multiprocessing import Pool
+from copy import deepcopy
+from typing import Literal
+from xai_components.serialization import save_params_npz
+from xai_components.settings import memory
+import os
+import json
+from xai_components.base_tvb import ComponentWithViewer
+from tvbextxircuits.logger.builder import get_logger
+
+LOGGER = get_logger(__name__)
+
+
+def cpp_worker(task):
+ from vbi import extract_features_df
+
+ (model_cls, base, theta_row, idx, nodewise, nn, fs, cfg, ts_key, t_key, return_data) = task
+ par_i = deepcopy(base)
+ for name, col in idx.items():
+ val = theta_row[col]
+ if name in nodewise:
+ par_i[name] = np.full(nn, val, dtype=float)
+ else:
+ par_i[name] = val
+ par_i.pop("dim", None) # temporary
+ data = model_cls(par_i).run()
+ ts = data[ts_key]
+ stat_vec = extract_features_df(ts=[ts], cfg=cfg, fs=fs,
+ n_workers=1, verbose=False).values
+
+ if return_data:
+ return stat_vec[0], data[t_key], data[ts_key]
+ else:
+ return stat_vec[0], None, None
+
+@xai_component(color='rgb(220, 5, 45)')
+class SimulationRunner(ComponentWithViewer):
+ backend: InArg[Literal['cupy', 'cpp']]
+ model: InArg[any]
+ theta: InArg[torch.Tensor]
+ theta_names: InArg[list]
+ cfg: InArg[dict]
+ num_workers: InArg[int]
+ time_series_key: InArg[dict]
+ output_dir: InArg[str]
+
+ stat_vec: OutArg[np.ndarray] # (N, F)
+
+ def __init__(self):
+ super().__init__()
+ self.backend.value = "cupy"
+ self.num_workers.value = 1
+
+ def execute(self, ctx):
+ # theta -> numpy
+ th = self.theta.value
+ theta_np = th.detach().cpu().numpy() if hasattr(th, "detach") else np.asarray(th)
+ theta_np = theta_np.astype(float, copy=False)
+ num_sim, num_params = theta_np.shape
+
+ if num_params != len(self.theta_names.value):
+ raise ValueError(f"Theta has {num_params} columns, but theta_names has {len(self.theta_names.value)}.")
+ idx = {par: index for index, par in enumerate(self.theta_names.value)}
+
+ fs = 1000.0 / float(self.model.value.dt)
+
+ # infer nn for broadcasting
+ nn = int(np.asarray(self.model.value.weights).shape[0])
+
+ nodewise = {"C0", "C1", "C2", "C3"} # temporary
+
+ model = self.model.value
+ ts_key = self.time_series_key.value['x']
+ t_key = self.time_series_key.value['t']
+
+ model_cls = model.__class__
+ base_par = deepcopy(model._par) # temporary
+ base_par["weights"] = np.array(base_par.get("weights"), dtype=np.float64)
+
+ if self.backend.value == "cpp":
+ tasks = []
+ for i in range(num_sim):
+ tasks.append((
+ model_cls,
+ base_par,
+ theta_np[i, :], # row for this sim
+ idx,
+ nodewise,
+ nn,
+ fs,
+ self.cfg.value,
+ ts_key, t_key,
+ i == 0
+ ))
+
+ with Pool(processes=self.num_workers.value) as pool:
+ results = pool.map(cpp_worker, tasks)
+
+ x = np.vstack([r[0] for r in results])
+
+ # Save data for timeseries viewer
+ resolved_par = self.build_resolved_par(base_par, theta_np, idx, nodewise, nn)
+ t0 = next((r[1] for r in results if r[1] is not None), None)
+ x0 = next((r[2] for r in results if r[2] is not None), None)
+
+ data = {ts_key: x0, t_key: t0}
+
+ elif self.backend.value == "cupy":
+ base_par["num_sim"] = num_sim
+
+ resolved_par = self.build_resolved_par(base_par, theta_np, idx, nodewise, nn)
+
+ data = simulate_cache(model_cls, resolved_par)
+
+ ts = data[ts_key]
+ if ts.ndim != 3:
+ raise ValueError(f"{self.backend.value} expected x=(time, nodes, nsim); got {ts.shape}")
+ ts = ts.transpose(2, 1, 0)
+ stat_vec = featurize_cache(ts, self.cfg.value, fs, int(self.num_workers.value), False)
+ x = stat_vec # (N, F)
+ else:
+ raise ValueError(f"{self.backend.value} backend not supported.")
+
+ self.stat_vec.value = x
+
+ self.persists_artifacts(self.output_dir.value, resolved_par, data, t_key, ts_key, self.theta_names.value)
+ LOGGER.info(f"Extracted features: {self.stat_vec.value}")
+
+ @staticmethod
+ def build_resolved_par(base_par, theta_np, idx, nodewise, nn):
+ for par, index in idx.items():
+ vals = theta_np[:, index]
+ if par in nodewise:
+ base_par[par] = np.tile(vals, (nn, 1))
+ else:
+ base_par[par] = vals
+
+ return base_par
+
+ @staticmethod
+ def persists_artifacts(output_dir, resolved_par, data, t_key, ts_key, theta_names):
+ # 1) model params
+ params_path = os.path.join(output_dir, "model_params.npz")
+ save_params_npz(resolved_par, params_path)
+
+ # 2) simulation data
+ data_path = os.path.join(output_dir, "simulation_data.npz")
+ np.savez(data_path, t=data[t_key], x=data[ts_key])
+
+ # 3) priors metadata
+ path = os.path.join(output_dir, "priors.json")
+ with open(path, "r", encoding="utf-8") as f:
+ priors_data = json.load(f)
+
+ priors_data["theta_names"] = theta_names
+ with open(path, "w", encoding="utf-8") as f:
+ json.dump(priors_data, f)
+
+
+@memory.cache
+def simulate_cache(model_class, resolved_par: dict) -> dict:
+ """
+ Cache the simulation results.
+ The cache key includes the model class and the model parameters
+ """
+ LOGGER.info(f"Resolved params: {resolved_par}")
+ model = model_class(resolved_par)
+ return model.run()
+
+@memory.cache(ignore=['n_workers', 'verbose'])
+def featurize_cache(ts, cfg: dict, fs: float, n_workers: int, verbose: bool) -> np.ndarray:
+ """
+ Cache the extracted features result.
+ The cache key includes the timeseries data, feature configuration dict (cfg) and sampling frequency (fs)
+ """
+ from vbi import extract_features
+
+ res = extract_features(ts=ts, cfg=cfg, fs=fs, n_workers=n_workers, verbose=verbose)
+ return res.values
diff --git a/xai_components/xai_vbi_train_posterior/__init__.py b/xai_components/xai_vbi_train_posterior/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/xai_components/xai_vbi_train_posterior/arguments/trainposterior.json b/xai_components/xai_vbi_train_posterior/arguments/trainposterior.json
new file mode 100644
index 00000000..f736966f
--- /dev/null
+++ b/xai_components/xai_vbi_train_posterior/arguments/trainposterior.json
@@ -0,0 +1,13 @@
+{
+ "description": "\n
Train Posterior
\n
Trains a posterior distribution from parameter samples (theta) and their corresponding simulated features (stat_vec) using an SBI method (SNPE/SNLE/SNRE). Before training, the feature matrix is standardized with StandardScaler to improve numerical stability.
\n
(N, D), where N is the number of samples/simulations and D is the dimension of the parameter space.\n",
+ "stat_vec" : "The feature samples, where n is the number of samples/simulations and d is the dimension of the feature space.\n",
+ "method" : "Inference method to train. Supported values are 'SNPE', 'SNLE', 'SNRE'. Default: 'SNPE'.\n",
+ "device" : "Device used for training. Supported values are 'cpu' and 'cuda'. Default: 'cpu'.\n",
+ "density_estimator" : "Neural density estimator used by the selected method. Supported values are 'maf' and 'nsf'. Default: 'maf'.\n",
+ "with_mean" : "Controls StandardScaler centering. When True, it centers the data before scaling.\n",
+ "with_std" : "Controls StandardScaler scaling. When True, it scales the data to unit variance (or equivalently, unit standard deviation).\n"
+ }
+}
\ No newline at end of file
diff --git a/xai_components/xai_vbi_train_posterior/train_posterior.py b/xai_components/xai_vbi_train_posterior/train_posterior.py
new file mode 100644
index 00000000..26bf1391
--- /dev/null
+++ b/xai_components/xai_vbi_train_posterior/train_posterior.py
@@ -0,0 +1,49 @@
+from xai_components.base import xai_component, Component, InArg, OutArg
+import numpy as np
+import torch
+from sklearn.preprocessing import StandardScaler
+import sbi.utils as utils
+
+@xai_component(color="rgb(220, 5, 45)")
+class TrainPosterior(Component):
+ prior: InArg[utils.BoxUniform]
+ theta: InArg[torch.Tensor]
+ stat_vec: InArg[np.ndarray]
+ method: InArg[str]
+ device: InArg[str]
+ density_estimator: InArg[str]
+ with_mean: InArg[bool]
+ with_std: InArg[bool]
+
+ posterior: OutArg[object]
+ X_scaled: OutArg[torch.Tensor]
+
+ def __init__(self):
+ super().__init__()
+ self.with_mean.value = True
+ self.with_std.value = True
+ self.method.value = "SNPE"
+ self.device.value = "cpu"
+ self.density_estimator.value = "maf"
+
+ def execute(self, ctx):
+ from vbi.sbi_inference import Inference
+
+ stat_vec_np = np.array(self.stat_vec.value)
+
+ scaler = StandardScaler(
+ with_mean=self.with_mean.value,
+ with_std=self.with_std.value
+ )
+ stat_vec_st = scaler.fit_transform(stat_vec_np)
+ stat_vec_st = torch.tensor(stat_vec_st , dtype=torch.float32)
+
+ # Train posterior
+ obj = Inference()
+ posterior = obj.train(self.theta.value, stat_vec_st, self.prior.value,
+ method=self.method.value, device=self.device.value,
+ density_estimator=self.density_estimator.value)
+
+ self.posterior.value = posterior
+ self.X_scaled.value = stat_vec_st
+