improve ExplicitAts API

Author: randriu

umbi/ats/__init__.py

@@ -16,7 +16,7 @@
     write,
 )
 from .entity_class import EntityClass
-from .explicit_ats import ExplicitAts, TimeType
+from .explicit_ats import Branch, Choice, ExplicitAts, TimeType
 from .model_info import ModelInfo
 from .variable_valuations import EntityClassValuations, EntityValuation, EntityValuations, Variable, VariableValuations
 
@@ -30,6 +30,8 @@
     # model_info
     "ModelInfo",
     # explicit_ats
+    "Branch",
+    "Choice",
     "TimeType",
     "ExplicitAts",
     # variable_valuations

umbi/ats/annotations.py

@@ -121,19 +121,19 @@ def observation_values(self) -> list[Scalar]:
     def player_values(self) -> list[Scalar]:
         return self.get_values_for(EntityClass.PLAYERS)
 
-    def set_state_values(self, values: list[Scalar]):
+    def set_state_values(self, values: Sequence[Scalar]):
         self.set_values_for(EntityClass.STATES, values)
 
-    def set_choice_values(self, values: list[Scalar]):
+    def set_choice_values(self, values: Sequence[Scalar]):
         self.set_values_for(EntityClass.CHOICES, values)
 
-    def set_branch_values(self, values: list[Scalar]):
+    def set_branch_values(self, values: Sequence[Scalar]):
         self.set_values_for(EntityClass.BRANCHES, values)
 
-    def set_observation_values(self, values: list[Scalar]):
+    def set_observation_values(self, values: Sequence[Scalar]):
         self.set_values_for(EntityClass.OBSERVATIONS, values)
 
-    def set_player_values(self, values: list[Scalar]):
+    def set_player_values(self, values: Sequence[Scalar]):
         self.set_values_for(EntityClass.PLAYERS, values)
 
     def unset_state_values(self):

umbi/ats/ats_to_umb.py

@@ -18,7 +18,8 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
     ats = ExplicitAts()
 
     ## index
-    # skip format_version, format_revision and file_data
+    # strip format_version, format_revision and file_data
+    # load umb.index.model_data
     md = umb.index.model_data
     if md is not None:
         ats.model_info = ModelInfo(
@@ -30,33 +31,51 @@ def explicit_umb_to_explicit_ats(umb: umbi.umb.ExplicitUmb) -> ExplicitAts:
             doi=md.doi,
             url=md.url,
         )
+
     # load index.transition_system
     ts = umb.index.transition_system
     ats.time = TimeType(ts.time)
     ats.num_players = ts.num_players
-    ats.num_states = ts.num_states
     ats.num_initial_states = ts.num_initial_states
-    ats.num_choices = ts.num_choices
     ats.num_choice_actions = ts.num_choice_actions
-    ats.num_branches = ts.num_branches
     ats.num_branch_actions = ts.num_branch_actions
     ats.player_to_name = ts.player_to_name
 
     ## values
     ats.state_is_initial = umb.state_is_initial
-    ats.state_to_choice = umb.state_to_choices
     ats.state_to_player = umb.state_to_player
 
     ats.state_is_markovian = umb.state_is_markovian
     ats.state_to_exit_rate = umb.state_to_exit_rate
 
-    ats.choice_to_branches = umb.choice_to_branches
-    ats.branch_to_target = umb.branch_to_target
-    ats.branch_to_probability = umb.branch_to_probability
+    for state_index in range(ts.num_states):
+        state = ats.add_state()
+        assert state == state_index
+
+    state_to_choices = (
+        umb.state_to_choices if umb.state_to_choices is not None else [s for s in range(ats.num_states + 1)]
+    )
+    if ts.num_choices > 0:
+        for state in range(ats.num_states):
+            for choice_idx in range(state_to_choices[state], state_to_choices[state + 1]):
+                choice = ats.add_state_choice(state=state)
+                if ts.num_branches > 0:
+                    assert umb.choice_to_branches is not None, "num_branches > 0 but choice_to_branches is None"
+                    for branch_idx in range(umb.choice_to_branches[choice_idx], umb.choice_to_branches[choice_idx + 1]):
+                        assert umb.branch_to_target is not None, "num_branches > 0 but branch_to_target is None"
+                        target = umb.branch_to_target[branch_idx]
+                        prob = umb.branch_to_probability[branch_idx] if umb.branch_to_probability is not None else None
+                        branch_action = (
+                            umb.branch_to_branch_action[branch_idx] if umb.branch_to_branch_action is not None else None
+                        )
+                        choice.add_branch(target=target, prob=prob, action=branch_action)
+                if ts.num_choice_actions > 0:
+                    assert umb.choice_to_choice_action is not None, (
+                        "num_choice_actions > 0 but choice_to_choice_action is None"
+                    )
+                    choice.action = umb.choice_to_choice_action[choice_idx]
 
-    ats.choice_to_choice_action = umb.choice_to_choice_action
     ats.choice_action_to_name = umb.choice_action_to_string
-    ats.branch_to_branch_action = umb.branch_to_branch_action
     ats.branch_action_to_name = umb.branch_action_to_string
 
     # load annotations
@@ -167,7 +186,6 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
     # warning: this does not copy the values, but directly uses the lists from ats.
     # this is fine as long as we don't modify them
     umb.state_is_initial = ats.state_is_initial
-    umb.state_to_choices = ats.state_to_choice
     umb.state_to_player = ats.state_to_player
 
     umb.state_is_markovian = ats.state_is_markovian
@@ -178,19 +196,40 @@ def explicit_ats_to_explicit_umb(ats: ExplicitAts) -> umbi.umb.ExplicitUmb:
         umb.index.transition_system.exit_rate_type = SizedType.for_type(target_type)
         umb.state_to_exit_rate = vector  # type: ignore
 
-    umb.choice_to_branches = ats.choice_to_branches
-    umb.branch_to_target = ats.branch_to_target
-
-    if ats.branch_to_probability is not None:
-        # promote
-        target_type, vector = umbi.datatypes.promote_scalars(ats.branch_to_probability)
-        assert isinstance(target_type, umbi.datatypes.NumericType), "branch probabilities must be numeric"
-        umb.index.transition_system.branch_probability_type = SizedType.for_type(target_type)
-        umb.branch_to_probability = vector  # type: ignore
+    umb.state_to_choices = []
+    num_choices = 0
+    for state in range(ats.num_states):
+        umb.state_to_choices.append(num_choices)
+        num_choices += ats.num_state_choices(state)
+    umb.state_to_choices.append(num_choices)
+
+    if ats.num_choices > 0:
+        if ats.num_choice_actions > 0:
+            umb.choice_to_choice_action = [choice.action for choice in ats.choices]  # type: ignore
+
+        branches = []
+        umb.choice_to_branches = []
+        for choice in ats.choices:
+            umb.choice_to_branches.append(len(branches))
+            branches.extend(choice.branches)
+        umb.choice_to_branches.append(len(branches))
+        if len(branches) > 0:
+            umb.branch_to_target = [branch.target for branch in branches]
+            if any(branch.prob is not None for branch in branches):
+                # has branch probabilities
+                branch_to_probability = [branch.prob if branch.prob is not None else 1 for branch in branches]
+                # promote
+                target_type, vector = umbi.datatypes.promote_scalars(branch_to_probability)
+                # assert isinstance(target_type, umbi.datatypes.NumericType), "branch probabilities must be numeric"
+                umb.index.transition_system.branch_probability_type = SizedType.for_type(target_type)
+                umb.branch_to_probability = vector  # type: ignore
+            if ats.num_branch_actions > 0:
+                branch_actions = [branch.action for branch in branches]
+                if any(action is None for action in branch_actions):
+                    raise ValueError("if num_branch_actions > 0, all branches must have an action")
+                umb.branch_to_branch_action = branch_actions  # type: ignore
 
-    umb.choice_to_choice_action = ats.choice_to_choice_action
     umb.choice_action_to_string = ats.choice_action_to_name
-    umb.branch_to_branch_action = ats.branch_to_branch_action
     umb.branch_action_to_string = ats.branch_action_to_name
 
     # add annotations

umbi/ats/examples/grid.py

@@ -1,18 +1,15 @@
-"""
-umbi demo: Create an ATS from a grid string.
-"""
-
-from __future__ import annotations
+"""umbi demo: Create an ExplicitAts from a grid string."""
 
 import logging
 from fractions import Fraction
 
-import umbi
+from ..explicit_ats import ExplicitAts, TimeType
+from ..variable_valuations import EntityClassValuations, EntityValuations
 
 logger = logging.getLogger(__name__)
 
 
-def grid_ats(grid: str) -> umbi.ats.ExplicitAts:
+def grid_ats(grid: str) -> ExplicitAts:
     """
     Create a simple ATS from a rectangular grid string.
 
@@ -32,8 +29,8 @@ def grid_ats(grid: str) -> umbi.ats.ExplicitAts:
         raise ValueError("The grid must be rectangular (all rows must have the same length).")
 
     # Create ATS
-    ats = umbi.ats.ExplicitAts()
-    ats.time = umbi.ats.TimeType.DISCRETE
+    ats = ExplicitAts()
+    ats.time = TimeType.DISCRETE
     ats.num_players = 1
 
     cell_to_state = {}
@@ -55,13 +52,12 @@ def grid_ats(grid: str) -> umbi.ats.ExplicitAts:
             elif c == "g":
                 goal_states.add(state)
 
-    ats.num_states = len(cell_to_state)
     ats.set_initial_states(list(initial_states))
 
     if ats.variable_valuations is None:
-        ats.variable_valuations = umbi.ats.EntityClassValuations()
+        ats.variable_valuations = EntityClassValuations()
     if not ats.variable_valuations.has_state_valuations:
-        ats.variable_valuations.set_state_valuations(umbi.ats.EntityValuations())
+        ats.variable_valuations.set_state_valuations(EntityValuations())
     state_valuations = ats.variable_valuations.state_valuations
     vx = state_valuations.add_variable("x")
     vy = state_valuations.add_variable("y")
@@ -76,41 +72,28 @@ def grid_ats(grid: str) -> umbi.ats.ExplicitAts:
         "right": (1, 0),
     }
 
-    ats.state_to_choice = []
-    ats.choice_to_branches = []
-    ats.branch_to_target = []
-    ats.branch_to_probability = []
-    ats.choice_to_choice_action = []
     ats.choice_action_to_name = list(direction_dxdy.keys())
     ats.num_choice_actions = len(ats.choice_action_to_name)
 
     for (x, y), state in cell_to_state.items():
-        ats.state_to_choice.append(len(ats.choice_to_choice_action))
-
+        ats.add_state()
         for direction, (dx, dy) in direction_dxdy.items():
             target = (x + dx, y + dy)
-            ats.choice_to_choice_action.append(ats.choice_action_to_name.index(direction))
-            ats.choice_to_branches.append(len(ats.branch_to_target))
+            choice = ats.add_state_choice(state=state)
+            choice.action = ats.choice_action_to_name.index(direction)
 
             if target in cell_to_state:
                 target_state = cell_to_state[target]
-                ats.branch_to_target.extend([target_state, state])
-                ats.branch_to_probability.extend([Fraction(9, 10), Fraction(1, 10)])
+                choice.add_branch(target=target_state, prob=Fraction(9, 10))
+                choice.add_branch(target=state, prob=Fraction(1, 10))
             else:
-                ats.branch_to_target.append(state)
-                ats.branch_to_probability.append(1)
-
-    ats.state_to_choice.append(len(ats.choice_to_choice_action))
-    ats.choice_to_branches.append(len(ats.branch_to_target))
-    ats.num_choices = len(ats.choice_to_choice_action)
-    ats.num_branches = len(ats.branch_to_target)
+                choice.add_branch(target=state)
 
-    goal_ap = umbi.ats.AtomicPropositionAnnotation(name="goal")
-    goal_ap.set_state_values([s in goal_states for s in range(ats.num_states)])
-    ats.add_ap_annotation(goal_ap)
+    annotation = ats.add_ap_annotation(name="goal")
+    annotation.set_state_values([s in goal_states for s in range(ats.num_states)])
 
-    step_cost_reward = umbi.ats.RewardAnnotation(name="step_cost")
-    step_cost_reward.set_choice_values([1 for _ in range(ats.num_choices)])
-    ats.add_reward_annotation(step_cost_reward)
+    annotation = ats.add_reward_annotation(name="step_cost")
+    annotation.set_choice_values([1 for _ in range(ats.num_choices)])
 
+    ats.validate()
     return ats

umbi/ats/examples/random_game.py

@@ -1,16 +1,13 @@
-#!/usr/bin/env python3
 """
-umbi demo: A random stochastic game generator based on Azeem et al. "Optimistic and Topological Value Iteration for
-Simple Stochastic Games".
+umbi demo: A random stochastic game generator based on
+Azeem et al. "Optimistic and Topological Value Iteration for Simple Stochastic Games".
 """
 
-from __future__ import annotations
-
 import logging
 import random
 import time
 
-import umbi
+from ..explicit_ats import ExplicitAts, TimeType
 
 logger = logging.getLogger(__name__)
 
@@ -176,7 +173,7 @@ def random_transition_function(
     return delta, actions_at_state, all_actions
 
 
-def random_game_ats(num_states: int, seed: int | None = None) -> umbi.ats.ExplicitAts:
+def random_game_ats(num_states: int, seed: int | None = None) -> ExplicitAts:
     """
     Generate a random stochastic game connected from initial state.
 
@@ -187,9 +184,8 @@ def random_game_ats(num_states: int, seed: int | None = None) -> umbi.ats.Explic
     delta, actions_at_state, all_actions = random_transition_function(num_states, seed)
 
     # create ATS
-    ats = umbi.ats.ExplicitAts(
-        time=umbi.ats.TimeType.STOCHASTIC,
-        num_states=num_states,
+    ats = ExplicitAts(
+        time=TimeType.STOCHASTIC,
         num_players=2,
     )
 
@@ -207,40 +203,25 @@ def random_game_ats(num_states: int, seed: int | None = None) -> umbi.ats.Explic
     ats.num_choice_actions = len(all_actions)
     ats.choice_action_to_name = all_actions
 
-    # count choices and branches
-    ats.num_choices = sum(len(actions) for actions in actions_at_state.values())
-    ats.num_branches = sum(
-        len(delta.get(s, {}).get(action, [])) for s in range(num_states) for action in actions_at_state[s]
-    )
-
     # build CSR structures
     logger.info("building CSR structures ...")
     build_start = time.time()
-    ats.state_to_choice = []
-    ats.choice_to_choice_action = []
-    ats.choice_to_branches = []
-    ats.branch_to_target = []
-    ats.branch_to_probability = []
 
     for s in range(num_states):
-        ats.state_to_choice.append(len(ats.choice_to_choice_action))
-
+        state = ats.add_state()
         for action in sorted(actions_at_state[s]):
-            ats.choice_to_choice_action.append(action)
-            ats.choice_to_branches.append(len(ats.branch_to_target))
+            choice = ats.add_state_choice(state=state)
+            choice.action = action
 
             # get all branches for this (s, action) pair
             branches = delta[s][action]
             branches.sort()  # sort by target state
 
             for target, prob in branches:
-                ats.branch_to_target.append(target)
-                ats.branch_to_probability.append(prob)
-
-    ats.state_to_choice.append(len(ats.choice_to_choice_action))
-    ats.choice_to_branches.append(len(ats.branch_to_target))
+                choice.add_branch(target=target, prob=prob)
 
     build_time = time.time() - build_start
+    ats.validate()
     logger.info(f"CSR building completed in {build_time:.3f}s")
     logger.info(f"generated random game: {num_states} states, {ats.num_choices} choices, {ats.num_branches} branches")
     logger.info(f"player0 states: {sum(is_player0_state)}, player1 states: {num_states - sum(is_player0_state)}")