Merge pull request #75 from the-virtual-brain/EBR-188

EBR-188: resolve OpenViewer issues after testing at JSC

Author: 915-Misan-Teodora

package.json

@@ -1,6 +1,6 @@
 {
   "name": "tvb-ext-xircuits",
-  "version": "3.0.0",
+  "version": "3.0.1",
   "description": "Jupyterlab extension for building TVB workflows in a visual and interactive manner",
   "keywords": [
     "jupyter",

pyproject.toml

@@ -68,7 +68,7 @@ full = [
     "tvb-gdist",
     "tvb-framework",
     "tvb-widgets>=1.0",
-    "vbi[inference]",
+    "vbi[inference]<=0.3",
     "torch",
     "sbi"
 ]
@@ -139,5 +139,6 @@ include = [
     "xai_components/",
     "tvbextxircuits/",
     "examples/",
+    "config/"
 ]
 artifacts = ["tvbextxircuits/labextension"]

tvbextxircuits/nb_generator.py

@@ -92,7 +92,10 @@ class NotebookFactory(object):
     def get_notebook_for_component(component_name, component_id, component_path, component_inputs, xircuits_id, xircuits_filename):
         component_class = determine_component_class(component_name, component_path)
 
-        if not (issubclass(component_class, ComponentWithWidget) or issubclass(component_class, ComponentWithViewer)):
+        if component_class is None:
+            return None
+
+        if not issubclass(component_class, (ComponentWithWidget, ComponentWithViewer)):
             return None
 
         if component_class.__name__.startswith('StoreResults'):
@@ -435,6 +438,14 @@ 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)
-    return component_class
+    try:
+        component_module = importlib.import_module(component_path.replace('/', '.')[:-3])
+        component_class = getattr(component_module, component_name)
+        return component_class
+    except ModuleNotFoundError as e:
+        LOGGER.error(f"Module not found while loading {component_path}: {e}")
+    except AttributeError:
+        LOGGER.error(f"Class {component_name} not found in module {component_path}.")
+
+    return None
+

xai_components/xai_vbi_sample_posterior/sample_posterior.py

@@ -1,6 +1,4 @@
 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
@@ -16,13 +14,14 @@ class SamplePosterior(ComponentWithViewer):
     obs_idx: InArg[int]
     output_dir: InArg[str]
 
-    samples: OutArg[torch.Tensor]
+    samples: OutArg[any]
 
     def __init__(self):
         super().__init__()
         self.obs_idx.value = 0
 
     def execute(self, ctx):
+        from vbi.sbi_inference import Inference
 
         x_idx_st = self.X_scaled.value[self.obs_idx.value,:]
 
@@ -38,6 +37,7 @@ def execute(self, ctx):
 
     @staticmethod
     def persists_artifacts(output_dir, samples):
+        import torch
         # Store the resulted samples for plotting
         path = os.path.join(output_dir, "samples.pt")
         torch.save(samples, path)

xai_components/xai_vbi_simulation_runner/simulation_runner.py

@@ -1,6 +1,5 @@
 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
@@ -40,7 +39,7 @@ def cpp_worker(task):
 class SimulationRunner(ComponentWithViewer):
     backend: InArg[Literal['cupy', 'cpp']]
     model: InArg[any]
-    theta: InArg[torch.Tensor]
+    theta: InArg[any]
     theta_names: InArg[list]
     cfg: InArg[dict]
     num_workers: InArg[int]