use custom clone of Vivarium class for more control during assembly

Author: jcschaff

src/simple_process_bigraph_runtime/environment/__init__.py

@@ -0,0 +1,245 @@
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+from bigraph_schema import is_schema_key, set_path, get_path  # type: ignore[import-untyped]
+from bigraph_schema.utilities import remove_path  # type: ignore[import-untyped]
+from process_bigraph import Composite, Process, Step, ProcessTypes  # type: ignore[import-untyped]
+
+
+class ProcessBigraphEnv:
+    core: ProcessTypes
+    composite: Composite
+    emitter_config: dict[str, Any]
+    emitter_paths: dict[str, tuple[str, ...]]
+
+    def __init__(self):
+        self.core = ProcessTypes()
+        self.composite = Composite(core=self.core)
+        self.emitter_config = {}
+        self.emitter_paths = {}
+
+    def add_process(self, name: str, process_id, config=None, inputs=None, outputs=None,
+                    path: tuple[str, ...] | None = None):
+        edge_type = "process"  # TODO -- does it matter if this is a step or a process?
+        config = config or {}
+        inputs = inputs or {}
+        outputs = outputs or {}
+        path = path or ()
+
+        # convert string paths to lists
+        # TODO -- make this into a separate path-parsing function
+        for ports in [inputs, outputs]:
+            for port, port_path in ports.items():
+                ports[port] = parse_path(port_path)
+
+        # make the process spec
+        state = {
+            name: {
+                "_type": edge_type,
+                "address": f"local:{process_id}",  # TODO -- only support local right now?
+                "config": config,
+                "inputs": inputs,
+                "outputs": outputs,
+            }
+        }
+
+        # nest the process in the composite at the given path
+        self.composite.merge({}, state, path)
+        self.reset_emitters()
+        self.reset_paths()
+
+    def add_step(self, name: str, process_id, config=None, inputs=None, outputs=None,
+                    path: tuple[str, ...] | None = None):
+        edge_type = "step"  # TODO -- does it matter if this is a step or a process?
+        config = config or {}
+        inputs = inputs or {}
+        outputs = outputs or {}
+        path = path or ()
+
+        # convert string paths to lists
+        # TODO -- make this into a separate path-parsing function
+        for ports in [inputs, outputs]:
+            for port, port_path in ports.items():
+                ports[port] = parse_path(port_path)
+
+        # make the process spec
+        state = {
+            name: {
+                "_type": edge_type,
+                "address": f"local:{process_id}",  # TODO -- only support local right now?
+                "config": config,
+                "inputs": inputs,
+                "outputs": outputs,
+            }
+        }
+
+        # nest the process in the composite at the given path
+        self.composite.merge({}, state, path)
+        self.reset_emitters()
+        self.reset_paths()
+
+    def reset_emitters(self) -> None:
+        for path, emitter in self.emitter_paths.items():
+            remove_path(self.composite.state, path)
+            self.add_emitter()
+#         for path, instance in self.composite.step_paths.items():
+#             if "emitter" in path:
+#                 remove_path(self.composite.state, path)
+#                 self.add_emitter()
+
+    def reset_paths(self):
+        self.composite.find_instance_paths(self.composite.state)
+        self.composite.build_step_network()
+
+    def add_emitter(self):
+        """
+        Add an emitter to the composite.
+        """
+
+        emitter_state = self._read_emitter_config()
+
+        # set the emitter at the path
+        path = tuple(self.emitter_config.get("path", ('emitter',)))
+        emitter_state = set_path({}, path, emitter_state)
+
+        self.composite.merge({}, emitter_state)
+
+        # TODO -- this is a hack to get the emitter to show up in the state
+        # TODO -- this should be done in the merge function
+        _, instance = self.core.slice(
+            self.composite.composition,
+            self.composite.state,
+            path)
+
+        self.emitter_paths[path] = instance
+        self.composite.step_paths[path] = instance
+
+        # rebuild the step network
+        self.composite.build_step_network()
+
+    def _read_emitter_config(self):
+        address = self.emitter_config.get("address", "local:ram-emitter")
+        config = self.emitter_config.get("config", {})
+        mode = self.emitter_config.get("mode", "all")
+
+        if mode == "all":
+            inputs = {}
+            for key in self.composite.state.keys():
+                if is_schema_key(key):  # skip schema keys
+                    continue
+                if self.core.inherits_from(self.composite.composition[key], "edge"):  # skip edges
+                    continue
+                inputs[key] = [self.emitter_config.get("inputs", {}).get(key, key)]
+
+        elif mode == "none":
+            inputs = self.emitter_config.get("emit", {})
+
+        elif mode == "bridge":
+            print("Warning: emitter bridge mode not implemented.")
+            inputs = {}
+
+        elif mode == "ports":
+            print("Warning: emitter ports mode not implemented.")
+            inputs = {}
+
+        if not "emit" in config:
+            config["emit"] = {
+                input: "any"
+                for input in inputs}
+
+        return {
+            "_type": "step",
+            "address": address,
+            "config": config,
+            "inputs": inputs}
+
+    def set_value(self, path: list[str], value: Any) -> None:
+        # TODO -- make this set the value in the composite using core
+        set_path(self.composite.state, path=path, value=value)
+        # self.composite[path] = value
+        # self.composite.set({}, value, path)
+
+    def get_type(self, type_id: str) -> dict[str, Any]:
+        return render_type(type_name_or_dict=self.core.access(type_id), core=self.core)
+
+    def make_document(self) -> dict[str, Any]:
+        serialized_state = self.composite.serialize_state()
+
+        # TODO fix RecursionError
+        # schema = self.core.representation(self.composite.composition)
+        schema = self.composite.serialize_schema()
+        # schema = self.composite.composition
+
+        return {
+            "state": serialized_state,
+            "composition": schema,
+        }
+
+    def save(self, json_file_path: Path):
+        document = self.make_document()
+        with open(json_file_path, "w") as json_file:
+            json.dump(document, json_file, indent=4)
+            print(f"Saved file: {json_file_path}")
+
+
+    def run(self, interval):
+        """
+        Run the simulation for a given interval.
+        """
+        if not self.emitter_paths:
+            self.add_emitter()
+
+        self.composite.run(interval)
+
+    def step(self):
+        """
+        Run the simulation for a single step.
+        """
+        self.composite.update({}, 0)
+
+
+
+def render_type(type_name_or_dict: str | dict[str, Any], core: ProcessTypes) -> dict[str, Any]:
+    type_dict: dict[str, Any] = core.access(type_name_or_dict)
+    rendered = {}
+    for key, value in type_dict.items():
+        if is_schema_key(key):
+            if key == '_description':
+                rendered['description'] = value
+            elif key == '_default':
+                rendered['default'] = core.default(type_dict)
+        elif isinstance(value, dict):
+            rendered[key] = render_type(value, core)
+    return rendered
+
+def parse_path(path: list[str] | str | dict[str, Any]) -> list[str] | dict[str, Any]:
+    if isinstance(path, dict):
+        return {
+            key: parse_path(value)
+            for key, value in path.items()}
+    elif isinstance(path, str):
+        if path.startswith('/'):
+            return path.split('/')[1:]
+        else:
+            return [path]
+    else:
+        return path
+
+
+class TypedStep(Step):
+    def __init__(self, config=None):
+        super().__init__(config)
+        self._inputs = self.inputs()
+        self._outputs = self.outputs()
+
+    def inputs(self) -> dict[str, Any]:
+        return {}
+
+
+    def outputs(self) -> dict[str, Any]:
+        return {}
+
+    def update(self, inputs: dict[str, Any]) -> dict[str, Any]:
+        raise NotImplementedError("Subclasses should implement this method.")

src/simple_process_bigraph_runtime/environment/process_bigraph_env.py

@@ -1,10 +1,11 @@
 from typing import Any, cast
 
 from process_bigraph_lang.dsl.model import Model, StoreDef, ProcessDef, Store, Type
-from vivarium import Vivarium  # type: ignore[import-untyped]
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
 
 
-def process_composite(model: Model, assembler: Vivarium):
+
+def process_composite(model: Model, assembler: ProcessBigraphEnv):
     for composite_def in model.compositeDefs:
 
         store_path_to_value_map: dict[str, Any] = {}
@@ -49,13 +50,27 @@ def process_composite(model: Model, assembler: Vivarium):
             if not process_def.python_path:
                 raise ValueError(f"Process definition {process_def.name} has no python path")
 
-            assembler.add_process(
-                name=process.name,
-                process_id=".".join(process_def.python_path.path),
-                config=process_config,
-                inputs=input_bindings,
-                outputs=output_bindings,
-            )
+            process_store_path = ()
+            if process_def.python_path.path[-1].endswith("Step"):
+                assembler.add_step(
+                    name=process.name,
+                    process_id=".".join(process_def.python_path.path),
+                    config=process_config,
+                    inputs=input_bindings,
+                    outputs=output_bindings,
+                    path=process_store_path
+                )
+            elif process_def.python_path.path[-1].endswith("Process"):
+                assembler.add_process(
+                    name=process.name,
+                    process_id=".".join(process_def.python_path.path),
+                    config=process_config,
+                    inputs=input_bindings,
+                    outputs=output_bindings,
+                    path=process_store_path
+                )
+            else:
+                raise ValueError(f"Process definition {process_def.name} has an invalid python path, expecting to end with Step or Process")
 
         for store_path_str in store_path_to_value_map:
             store_path = store_path_str.split("::")

src/simple_process_bigraph_runtime/generation/composite_generator.py

@@ -1,11 +1,11 @@
 from process_bigraph_lang.dsl.model import ProcessDef
-from vivarium import Vivarium  # type: ignore[import-untyped]
-from vivarium.vivarium import VivariumTypes  # type: ignore[import-untyped]
 
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
 
-def register_process_defs(assembler: Vivarium, process_defs: list[ProcessDef]) -> None:
+
+def register_process_defs(assembler: ProcessBigraphEnv, process_defs: list[ProcessDef]) -> None:
     # note: does not actually dynamically build processes...yet
-    registry: VivariumTypes = assembler.core
+    registry = assembler.core
     for process_def in process_defs:
         if not process_def.python_path:
             raise ValueError(f"Process definition {process_def.name} has no python path")

src/simple_process_bigraph_runtime/generation/process_generator.py

@@ -1,9 +1,10 @@
 from process_bigraph_lang.dsl.model import Type
-from vivarium import Vivarium  # type: ignore[import-untyped]
 
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
 
-def register_types(assembler: Vivarium, types_to_register: list[Type]):
+
+def register_types(assembler: ProcessBigraphEnv, types_to_register: list[Type]):
     for type_to_register in types_to_register:
-        if 0 == assembler.get_type(type_to_register.name).size:
+        if assembler.get_type(type_to_register.name) == {}:
             raise ValueError(f"Type {type_to_register.name} is an unknown type")
         # in the future, we'd attempt to generate a new type

src/simple_process_bigraph_runtime/generation/type_generator.py

@@ -1,7 +1,8 @@
 from process_bigraph_lang.dsl.model import Unit
-from vivarium import Vivarium  # type: ignore[import-untyped]
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
 
-def register_units(assembler: Vivarium, units_to_register: list[Unit]) -> None:
+
+def register_units(assembler: ProcessBigraphEnv, units_to_register: list[Unit]) -> None:
     for unit_to_register in units_to_register:
         pass
         # if 0 == assembler.get_type(unit_to_register.name).size:

src/simple_process_bigraph_runtime/generation/unit_generator.py

@@ -12,6 +12,7 @@
 import simple_process_bigraph_runtime.registry.spatio_flux_library as spatioflux
 import simple_process_bigraph_runtime.registry.toy_library as toy
 from simple_process_bigraph_runtime import dsl_adapter
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
 from simple_process_bigraph_runtime.generation.composite_generator import process_composite
 from simple_process_bigraph_runtime.generation.process_generator import register_process_defs
 from simple_process_bigraph_runtime.generation.type_generator import register_types
@@ -70,9 +71,13 @@ def validate_pb_absolute_path(absolute_path: str) -> str:
     return absolute_path
 
 def performConversion(ast_model: Model) -> tuple[Composite, dict, ProcessTypes]:
-    assembler = Vivarium()
-    spatioflux.apply_to_vivarium(assembler)
-    toy.apply_to_vivarium(assembler)
+    assembler = ProcessBigraphEnv()
+
+    # register the libraries
+    spatioflux.register(assembler)
+    toy.register(assembler)
+
+    # register the types, units, and process definitions from the model
     register_types(assembler, ast_model.types)
     register_units(assembler, ast_model.units)
     register_process_defs(assembler, ast_model.processDefs)

src/simple_process_bigraph_runtime/main.py

@@ -2,6 +2,8 @@
 from spatio_flux import processes  # type: ignore[import-untyped]
 from vivarium import Vivarium  # type: ignore[import-untyped]
 
+from simple_process_bigraph_runtime.environment.process_bigraph_env import ProcessBigraphEnv
+
 
 def apply_non_negative(schema, current, update, top_schema, top_state, path, core):
     new_value = current + update
@@ -48,8 +50,8 @@ def apply_non_negative(schema, current, update, top_schema, top_state, path, cor
 
 PROCESS_DICT = processes.PROCESS_DICT.copy()
 
-def apply_to_vivarium(vivarium: Vivarium) -> None:
+def register(assembler: ProcessBigraphEnv) -> None:
     for type_name, type_schema in TYPES_DICT.items():
-        vivarium.core.register(type_name, type_schema)
+        assembler.core.register(type_name, type_schema)
     for process_name, process in PROCESS_DICT.items():
-        vivarium.core.register_process(process_name, process)
+        assembler.core.register_process(process_name, process)