fix workflow, examples

Author: randriu

.github/workflows/bump.yml

@@ -23,7 +23,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.11'
+          python-version: '3.12'
 
       - name: install bump-my-version
         run: |

.github/workflows/release.yml

@@ -50,7 +50,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.11'
+          python-version: '3.12'
 
       - name: install uv
         uses: astral-sh/setup-uv@v3

.github/workflows/test-build.yml

@@ -23,10 +23,10 @@ jobs:
           args: "--version"
 
       - name: lint code
-        run: ruff check --output-format=github --target-version=py311
+        run: ruff check --output-format=github --target-version=py312
 
       - name: check code formatting
-        run: ruff format --diff --target-version=py311
+        run: ruff format --diff --target-version=py312
         continue-on-error: true # Remove this once the project is formatted
 
       - name: install uv

README.md

@@ -15,12 +15,12 @@ pip install umbi
 
 ## Quick start
 
-A short example where we read a umbfile into an [`ExplicitAts`](umbi/ats/explicit_ats.py) object, modify initial states, and write it back:
+A short example where we read a umbfile into an [`SimpleAts`](umbi/ats/explicit_ats.py) object, modify initial states, and write it back:
 
 ```python
 import umbi
 ats = umbi.ats.read("in.umb")
-ats.set_initial_states([ats.num_states - 1])
+ats.initial_states = [ats.num_states - 1]
 umbi.ats.write(ats, "out.umb")
 ```
 
@@ -34,15 +34,15 @@ More examples can be found in the [examples](./examples) folder.
 
 **[`ExplicitUmb`](umbi/umb/explicit_umb.py)** - in-memory representation of a typical umbfile. Attributes are standard Python objects (lists, dicts, dataclasses) providing a deserialized view of the file contents.
 
-**[`ExplicitAts`](umbi/ats/explicit_ats.py)** - format-agnostic abstraction for annotated transition systems (states, transitions, annotations). Recommended for most use cases: easiest to use programmatically and remains stable across UMB format changes. See [umbi.ats.examples](./umbi/ats/examples/) for usage examples.
+**[`SimpleAts`](umbi/ats/simple_ats.py)** - format-agnostic abstraction for annotated transition systems (states, transitions, annotations, etc.). Recommended for most use cases: easiest to use programmatically and remains stable across UMB format changes. See [umbi.ats.examples](./umbi/ats/examples/) for usage examples.
 
 ## CLI
 
 `umbi` provides a basic CLI for umbfile manipulation.
 
 **Options:**
 - `--import-umb <path>` - import umbfile as `ExplicitUmb`
-- `--import-ats <path>` - import umbfile as `ExplicitAts`
+- `--import-ats <path>` - import umbfile as `SimpleAts`
 - `--export <path>` - export to umbfile (requires `--import-umb` or `--import-ats`)
 - `--log-level <LEVEL>` - set logging level: `DEBUG`, `INFO` (default), `WARNING`, `ERROR`, `CRITICAL`
 

examples/demo/ats.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Demonstration of umbfile manipulation via ExplicitAts."""
+"""Demonstration of umbfile manipulation via SimpleAts."""
 
 import pathlib
 
@@ -31,11 +31,11 @@ def main(input: pathlib.Path, output: pathlib.Path):
 
     print("Adding custom rewards...")
     annotation_name = "state_index"
-    reward_annotation = ats.add_reward_annotation(
+    reward_annotation = ats.new_reward_annotation(
         name=annotation_name,
         description="Reward annotation for states, equal to the state index",
     )
-    reward_annotation.set_state_values(list(ats.states))
+    reward_annotation.state_values = list(ats.states)
 
     print(f"Writing modified ATS to {output}...")
     umbi.ats.write(ats, output)

examples/misc/ats_smg_to_prism.py

@@ -8,6 +8,7 @@
 import pathlib
 
 import umbi
+import umbi.ats
 
 log = logging.getLogger(__name__)
 

examples/misc/value_iteration.py

@@ -11,7 +11,7 @@
 
 
 def reachability_vi(
-    model: umbi.ats.ExplicitAts,
+    model: umbi.ats.SimpleAts,
     goal_states: set[int],
     maximizing_players: set[int],
     stopping_criterion=1e-6,

umbi/ats/choices_mixin.py

@@ -1,10 +1,9 @@
 import logging
-from collections.abc import Sequence
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass, field
 
 from umbi.datatypes import Numeric, is_numeric_a_probability
 
-from .base_mixin import BaseMixin, CompositeMixin
 from .entity_space import EntityMapping
 from .entity_space_mixins import HasBranchSpace, HasChoiceSpace, HasStateSpace
 
@@ -155,12 +154,24 @@ def validate(self) -> None:
             # TODO check that probabilities sum to 1
         super().validate()
 
-    def new_state_choice(self, state: int) -> int:
-        """Add a choice to the given state."""
+    def new_state_choice(
+        self,
+        state: int,
+        targets: Iterable[int] | None = None,
+        target_prob: Callable[[int], Numeric | None] | None = None,
+    ) -> int:
+        """Add a choice to the given state.
+        :param state: source state for the new choice
+        :param targets: optional iterable of target states for the branches of the new choice
+        :param target_prob: optional function mapping target states to branch probabilities
+        :return: the new choice index
+        """
         # self.state_space._check_entity(state)
         choice = self._new_choice()
         self.state_to_choices[state].append(choice)
         self._choice_to_state[choice] = state
+        if targets is not None and target_prob is not None:
+            self.new_choice_branches(choice, targets=targets, target_prob=target_prob)
         return choice
 
     def remove_choice(self, choice: int) -> list[int]:
@@ -196,6 +207,22 @@ def new_choice_branch(self, choice: int, target: int | None = None, prob: Numeri
             self._branch_to_probability[new_branch] = prob
         return new_branch
 
+    def new_choice_branches(
+        self, choice: int, targets: Iterable[int], target_prob: Callable[[int], Numeric | None]
+    ) -> list[int]:
+        """Add branches to the given choice for multiple targets.
+        :param choice: choice to add branches to
+        :param targets: target states for the new branches
+        :param target_prob: function mapping target states to branch probabilities
+        :return: list of new branches added
+        """
+        new_branches = []
+        for target in targets:
+            prob = target_prob(target)
+            branch = self.new_choice_branch(choice, target=target, prob=prob)
+            new_branches.append(branch)
+        return new_branches
+
     def remove_branch(self, branch: int) -> None:
         self._branch_space.check_entity(branch)
         choice = self._branch_to_choice[branch]