improve ExplicitAts API

- ExplicitAts: improve API for state, choice and branch manipulation
- ExplicitAts: add utilities for state reordering
- readme: remove tag badge
- UmbIndex: more validations

Author: randriu

README.md

@@ -3,7 +3,6 @@
 [![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue)](https://www.python.org/)
 [![License MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
 [![PyPI](https://img.shields.io/pypi/v/umbi)](https://pypi.org/project/umbi/)
-[![Latest Tag](https://img.shields.io/github/tag/pmc-tools/umbi.svg)](https://github.com/pmc-tools/umbi/tags)
 [![Build Status](https://img.shields.io/github/actions/workflow/status/pmc-tools/umbi/test-build.yml)](https://github.com/pmc-tools/umbi/actions)
 
 Library for input/output of annotated transition systems (ATSs) in a *unified Markov binary (UMB)* format. See the [format specification](https://github.com/pmc-tools/umb/) for details.

tests/binary/test_sized_type.py

@@ -67,10 +67,10 @@ class TestSizedTypeValidation:
 
     def test_positive_size_required(self):
         """Test that size must be positive."""
-        with pytest.raises(AssertionError):
+        with pytest.raises(ValueError):
             SizedType(NumericPrimitiveType.UINT, 0)
 
-        with pytest.raises(AssertionError):
+        with pytest.raises(ValueError):
             SizedType(NumericPrimitiveType.UINT, -1)
 
     def test_string_must_be_64_bits(self):
@@ -147,3 +147,50 @@ def test_double_default_size(self):
         """Test default size for DOUBLE."""
         sized_type = SizedType.for_type(NumericPrimitiveType.DOUBLE)
         assert sized_type.size_bits == 64
+
+
+class TestSizedTypeNumericProperties:
+    """Test SizedType numeric classification properties."""
+
+    def test_is_numeric_for_numeric_types(self):
+        """Test is_numeric returns True for numeric types."""
+        assert SizedType(NumericPrimitiveType.INT, 32).is_numeric
+        assert SizedType(NumericPrimitiveType.UINT, 32).is_numeric
+        assert SizedType(NumericPrimitiveType.DOUBLE, 64).is_numeric
+        assert SizedType(NumericPrimitiveType.RATIONAL, 128).is_numeric
+
+    def test_is_numeric_for_primitive_types(self):
+        """Test is_numeric returns False for primitive types."""
+        assert not SizedType(PrimitiveType.BOOL, 1).is_numeric
+        assert not SizedType(PrimitiveType.STRING, 64).is_numeric
+
+    def test_is_discrete_for_discrete_types(self):
+        """Test is_discrete returns True for discrete numeric types."""
+        assert SizedType(NumericPrimitiveType.INT, 32).is_discrete
+        assert SizedType(NumericPrimitiveType.UINT, 32).is_discrete
+        assert SizedType(NumericPrimitiveType.INT, 64).is_discrete
+
+    def test_is_discrete_for_continuous_types(self):
+        """Test is_discrete returns False for continuous numeric types."""
+        assert not SizedType(NumericPrimitiveType.DOUBLE, 64).is_discrete
+        assert not SizedType(NumericPrimitiveType.RATIONAL, 128).is_discrete
+
+    def test_is_discrete_for_primitive_types(self):
+        """Test is_discrete returns False for primitive types."""
+        assert not SizedType(PrimitiveType.BOOL, 1).is_discrete
+        assert not SizedType(PrimitiveType.STRING, 64).is_discrete
+
+    def test_is_continuous_for_continuous_types(self):
+        """Test is_continuous returns True for continuous numeric types."""
+        assert SizedType(NumericPrimitiveType.DOUBLE, 64).is_continuous
+        assert SizedType(NumericPrimitiveType.RATIONAL, 128).is_continuous
+
+    def test_is_continuous_for_discrete_types(self):
+        """Test is_continuous returns False for discrete numeric types."""
+        assert not SizedType(NumericPrimitiveType.INT, 32).is_continuous
+        assert not SizedType(NumericPrimitiveType.UINT, 32).is_continuous
+
+    def test_is_continuous_for_primitive_types(self):
+        """Test is_continuous returns False for primitive types."""
+        assert not SizedType(PrimitiveType.BOOL, 1).is_continuous
+        assert not SizedType(PrimitiveType.STRING, 64).is_continuous

tests/datatypes/test_interval.py

@@ -78,6 +78,44 @@ def test_interval_equal_bounds(self):
         assert interval.left == 5
         assert interval.right == 5
 
+
+class TestIntervalTypeClassification:
+    """Test IntervalType classification properties."""
+
+    def test_is_discrete_for_int_interval(self):
+        """Test that int-interval is discrete."""
+        it = IntervalType(NumericPrimitiveType.INT)
+        assert it.is_discrete
+
+    def test_is_discrete_for_uint_interval(self):
+        """Test that uint-interval is discrete."""
+        it = IntervalType(NumericPrimitiveType.UINT)
+        assert it.is_discrete
+
+    def test_is_discrete_for_continuous_intervals(self):
+        """Test that double-interval and rational-interval are not discrete."""
+        double_it = IntervalType(NumericPrimitiveType.DOUBLE)
+        rational_it = IntervalType(NumericPrimitiveType.RATIONAL)
+        assert not double_it.is_discrete
+        assert not rational_it.is_discrete
+
+    def test_is_continuous_for_double_interval(self):
+        """Test that double-interval is continuous."""
+        it = IntervalType(NumericPrimitiveType.DOUBLE)
+        assert it.is_continuous
+
+    def test_is_continuous_for_rational_interval(self):
+        """Test that rational-interval is continuous."""
+        it = IntervalType(NumericPrimitiveType.RATIONAL)
+        assert it.is_continuous
+
+    def test_is_continuous_for_discrete_intervals(self):
+        """Test that int-interval and uint-interval are not continuous."""
+        int_it = IntervalType(NumericPrimitiveType.INT)
+        uint_it = IntervalType(NumericPrimitiveType.UINT)
+        assert not int_it.is_continuous
+        assert not uint_it.is_continuous
+
     def test_interval_str_repr(self):
         """Test Interval string representation."""
         interval = Interval(1, 5)

tests/datatypes/test_numeric_primitive.py

@@ -52,6 +52,36 @@ def test_bool_not_numeric_primitive(self):
         assert numeric_primitive_type_of(True) == NumericPrimitiveType.INT
 
 
+class TestNumericPrimitiveTypeClassification:
+    """Test numeric type classification properties."""
+
+    def test_is_discrete_for_int(self):
+        """Test that INT is discrete."""
+        assert NumericPrimitiveType.INT.is_discrete
+
+    def test_is_discrete_for_uint(self):
+        """Test that UINT is discrete."""
+        assert NumericPrimitiveType.UINT.is_discrete
+
+    def test_is_discrete_for_continuous_types(self):
+        """Test that DOUBLE and RATIONAL are not discrete."""
+        assert not NumericPrimitiveType.DOUBLE.is_discrete
+        assert not NumericPrimitiveType.RATIONAL.is_discrete
+
+    def test_is_continuous_for_double(self):
+        """Test that DOUBLE is continuous."""
+        assert NumericPrimitiveType.DOUBLE.is_continuous
+
+    def test_is_continuous_for_rational(self):
+        """Test that RATIONAL is continuous."""
+        assert NumericPrimitiveType.RATIONAL.is_continuous
+
+    def test_is_continuous_for_discrete_types(self):
+        """Test that INT and UINT are not continuous."""
+        assert not NumericPrimitiveType.INT.is_continuous
+        assert not NumericPrimitiveType.UINT.is_continuous
+
+
 class TestNumericPrimitivePromotionType:
     """Test numeric_primitive_promotion_type function."""
 

umbi/__main__.py

@@ -3,6 +3,7 @@
 import click
 
 import umbi
+import umbi.umb
 
 logger = logging.getLogger(__name__)
 
@@ -36,9 +37,9 @@ def main(log_level, import_umb, import_ats, export):
             raise ValueError("--export specified, but nothing to export")
         if umb is not None and ats is not None:
             raise ValueError("cannot specify both --import-umb and --import-ats when using --export")
-        if import_umb is not None:
+        if umb is not None:
             umbi.umb.write(umb, export)
-        if import_ats is not None:
+        if ats is not None:
             umbi.ats.write(ats, export)
 
 

umbi/ats/__init__.py

@@ -16,9 +16,15 @@
     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
+from .variable_valuations import (
+    EntityClassValuations,
+    EntityValuation,
+    EntityValuations,
+    EntityVariableValuations,
+    Variable,
+)
 
 __all__ = [
     "examples",
@@ -30,11 +36,13 @@
     # model_info
     "ModelInfo",
     # explicit_ats
+    "Branch",
+    "Choice",
     "TimeType",
     "ExplicitAts",
     # variable_valuations
+    "EntityVariableValuations",
     "Variable",
-    "VariableValuations",
     "EntityValuation",
     "EntityValuations",
     "EntityClassValuations",

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):
@@ -172,6 +172,18 @@ def validate(self):
         """Validate the annotation data."""
         return
 
+    def _reorder_entities(self, entity_class: EntityClass, new_order: list[int]) -> None:
+        """Reorder the values for the given entity class according to the new order.
+
+        :param entity_class: the entity class for which to reorder values
+        :param new_order: a list of indices representing the new order of values. Must be a permutation of the full list of entity indices.
+        """
+        if not self.has_values_for(entity_class):
+            return
+        values = self._entity_class_to_values[entity_class]
+        # we don't check the order
+        self._entity_class_to_values[entity_class] = [values[i] for i in new_order]
+
 
 class RewardAnnotation(Annotation):
     """Reward annotation is a numeric annotation that can be associated with states, choices, or branches."""

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 ats.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 ats.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