pgc state object can now contain arbitrary attributes

Author: PimLeerkes

stormvogel/model.py

@@ -205,7 +205,7 @@ class Action:
         labels: The labels of this action. Corresponds to Storm labels.
     """
 
-    name: str
+    name: str  # TODO name is stormpy choice label or we don't need a name at all?
     labels: frozenset[str]
 
     def __str__(self):

stormvogel/pgc.py

@@ -3,27 +3,36 @@
 from typing import Callable
 
 
+@dataclass
+class Action:
+    """pgc action object. Contains a list of labels"""
+
+    labels: list[str]
+
+
 @dataclass
 class State:
-    f: dict[str, int]
+    """pgc state object. Can contain any number of any type of arguments"""
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __repr__(self):
+        return f"State({self.__dict__})"
 
     def __hash__(self):
-        return hash(self.x)
+        return hash(self.__dict__)
 
     def __eq__(self, other):
         if isinstance(other, State):
-            return self.f == other.f
+            return self.__dict__ == other.__dict__
         return False
 
 
-@dataclass
-class Action:
-    labels: list[str]
-
-
 def build_pgc(
     delta,  # Callable[[State, Action], list[tuple[float, State]]],
-    initial_state_pgc: State,
+    initial_state_pgc: State,  # TODO rewards function, label function
     available_actions: Callable[[State], list[Action]] | None = None,
     modeltype: stormvogel.model.ModelType = stormvogel.model.ModelType.MDP,
 ) -> stormvogel.model.Model:
@@ -46,7 +55,9 @@ def build_pgc(
 
     # we create the model with the given type and initial state
     model.new_state(
-        labels=["init"], features=initial_state_pgc.f, name=str(initial_state_pgc.f)
+        labels=["init"],
+        features=initial_state_pgc.__dict__,
+        name=str(initial_state_pgc.__dict__),
     )
 
     # we continue calling delta and adding new states until no states are
@@ -82,14 +93,15 @@ def build_pgc(
                     if tuple[1] not in states_seen:
                         states_seen.append(tuple[1])
                         new_state = model.new_state(
-                            name=str(tuple[1].f), features=tuple[1].f
+                            name=str(tuple[1].__dict__), features=tuple[1].__dict__
                         )
                         branch.append((tuple[0], new_state))
                         states_to_be_visited.append(tuple[1])
                     else:
-                        # TODO what if there are multiple states with the same label? use names?
+                        # print(tuple[1].__dict__)
+                        # print(model.states)
                         branch.append(
-                            (tuple[0], model.get_state_by_name(str(tuple[1].f)))
+                            (tuple[0], model.get_state_by_name(str(tuple[1].__dict__)))
                         )
                 if branch != []:
                     transition[stormvogel_action] = stormvogel.model.Branch(branch)
@@ -101,19 +113,20 @@ def build_pgc(
                 if tuple[1] not in states_seen:
                     states_seen.append(tuple[1])
                     new_state = model.new_state(
-                        name=str(tuple[1].f), features=tuple[1].f
+                        name=str(tuple[1].__dict__), features=tuple[1].__dict__
                     )
 
                     branch.append((tuple[0], new_state))
                     states_to_be_visited.append(tuple[1])
                 else:
-                    # TODO what if there are multiple states with the same label? use names?
-                    branch.append((tuple[0], model.get_state_by_name(str(tuple[1].f))))
+                    branch.append(
+                        (tuple[0], model.get_state_by_name(str(tuple[1].__dict__)))
+                    )
                 if branch != []:
                     transition[stormvogel.model.EmptyAction] = stormvogel.model.Branch(
                         branch
                     )
-        s = model.get_state_by_name(str(state.f))
+        s = model.get_state_by_name(str(state.__dict__))
         assert s is not None
         model.add_transitions(
             s,

tests/test_pgc.py

@@ -1,15 +1,14 @@
 from stormvogel import pgc
 import stormvogel.model
 import math
-import stormpy
 import stormvogel.mapping
 
 
 def test_pgc_mdp():
     # we build the model with pgc:
     N = 2
     p = 0.5
-    initial_state = pgc.State({"x": math.floor(N / 2)})
+    initial_state = pgc.State(x=math.floor(N / 2))
 
     left = pgc.Action(["left"])
     right = pgc.Action(["right"])
@@ -19,23 +18,21 @@ def available_actions(s: pgc.State):
 
     def delta(s: pgc.State, action: pgc.Action):
         if action == left:
-            print(s.f)
-            print(s.f["x"])
             return (
                 [
-                    (p, pgc.State({"x": s.f["x"] + 1})),
-                    (1 - p, pgc.State({"x": s.f["x"]})),
+                    (p, pgc.State(x=s.x + 1)),
+                    (1 - p, pgc.State(x=s.x)),
                 ]
-                if s.f["x"] < N
+                if s.x < N
                 else []
             )
         elif action == right:
             return (
                 [
-                    (p, pgc.State({"x": s.f["x"] - 1})),
-                    (1 - p, pgc.State({"x": s.f["x"]})),
+                    (p, pgc.State(x=s.x - 1)),
+                    (1 - p, pgc.State(x=s.x)),
                 ]
-                if s.f["x"] > 0
+                if s.x > 0
                 else []
             )
 
@@ -63,41 +60,38 @@ def delta(s: pgc.State, action: pgc.Action):
     model.add_transitions(state2, stormvogel.model.Transition({right: branch2}))
     model.add_transitions(state0, stormvogel.model.Transition({left: branch0}))
 
-    print(model)
-    print(pgc_model)
-
     assert model == pgc_model
 
 
 def test_pgc_dtmc():
     # we build the model with pgc:
     p = 0.5
-    initial_state = pgc.State({"s": 0})
+    initial_state = pgc.State(s=0)
 
     def delta(s: pgc.State):
-        match s.f["s"]:
+        match s.s:
             case 0:
-                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 2}))]
+                return [(p, pgc.State(s=1)), (1 - p, pgc.State(s=2))]
             case 1:
-                return [(p, pgc.State({"s": 3})), (1 - p, pgc.State({"s": 4}))]
+                return [(p, pgc.State(s=3)), (1 - p, pgc.State(s=4))]
             case 2:
-                return [(p, pgc.State({"s": 5})), (1 - p, pgc.State({"s": 6}))]
+                return [(p, pgc.State(s=5)), (1 - p, pgc.State(s=6))]
             case 3:
-                return [(p, pgc.State({"s": 1})), (1 - p, pgc.State({"s": 7, "d": 1}))]
+                return [(p, pgc.State(s=1)), (1 - p, pgc.State(s=7, d=1))]
             case 4:
                 return [
-                    (p, pgc.State({"s": 7, "d": 2})),
-                    (1 - p, pgc.State({"s": 7, "d": 3})),
+                    (p, pgc.State(s=7, d=2)),
+                    (1 - p, pgc.State(s=7, d=3)),
                 ]
             case 5:
                 return [
-                    (p, pgc.State({"s": 7, "d": 4})),
-                    (1 - p, pgc.State({"s": 7, "d": 5})),
+                    (p, pgc.State(s=7, d=4)),
+                    (1 - p, pgc.State(s=7, d=5)),
                 ]
             case 6:
-                return [(p, pgc.State({"s": 2})), (1 - p, pgc.State({"s": 7, "d": 6}))]
+                return [(p, pgc.State(s=2)), (1 - p, pgc.State(s=7, d=6))]
             case 7:
-                return [(1, pgc.State({"s": 7}))]
+                return [(1, pgc.State(s=7))]
 
     pgc_model = stormvogel.pgc.build_pgc(
         delta=delta,
@@ -106,12 +100,65 @@ def delta(s: pgc.State):
     )
 
     # we build the model in the regular way:
-    path = stormpy.examples.files.prism_dtmc_die
-    prism_program = stormpy.parse_prism_program(path)
-    formula_str = "P=? [F s=7 & d=2]"
-    properties = stormpy.parse_properties(formula_str, prism_program)
-    model = stormpy.build_model(prism_program, properties)
-    print(dir(model.states[0]))
-    stormvogel_model = stormvogel.mapping.stormpy_to_stormvogel(model)
-    print(pgc_model)
-    print(stormvogel_model)
+    model = stormvogel.model.new_dtmc()
+    model.states[0].features = {"s": 0}
+    model.set_transitions(
+        model.get_initial_state(),
+        [
+            (1 / 2, model.new_state(features={"s": 1})),
+            (1 / 2, model.new_state(features={"s": 2})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(1),
+        [
+            (1 / 2, model.new_state(features={"s": 3})),
+            (1 / 2, model.new_state(features={"s": 4})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(2),
+        [
+            (1 / 2, model.new_state(features={"s": 5})),
+            (1 / 2, model.new_state(features={"s": 6})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(3),
+        [
+            (1 / 2, model.get_state_by_id(1)),
+            (1 / 2, model.new_state(features={"s": 7, "d": 1})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(4),
+        [
+            (1 / 2, model.new_state(features={"s": 7, "d": 2})),
+            (1 / 2, model.new_state(features={"s": 7, "d": 3})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(5),
+        [
+            (1 / 2, model.new_state(features={"s": 7, "d": 4})),
+            (1 / 2, model.new_state(features={"s": 7, "d": 5})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(6),
+        [
+            (1 / 2, model.get_state_by_id(2)),
+            (1 / 2, model.new_state(features={"s": 7, "d": 6})),
+        ],
+    )
+    model.set_transitions(
+        model.get_state_by_id(7), [(1, model.new_state(features={"s": 7}))]
+    )
+    model.set_transitions(model.get_state_by_id(8), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(9), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(10), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(11), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(12), [(1, model.get_state_by_id(13))])
+    model.set_transitions(model.get_state_by_id(13), [(1, model.get_state_by_id(13))])
+
+    assert pgc_model == model