better valuations API

Author: randriu

umbi/ats/examples/grid.py

@@ -71,9 +71,7 @@ def grid(grid: str) -> ExplicitAts:
         raise ValueError("The grid must be rectangular (all rows must have the same length).")
 
     # Create ATS
-    ats = ExplicitAts()
-    ats.time = TimeType.DISCRETE
-    ats.num_players = 1
+    ats = ExplicitAts(time=TimeType.DISCRETE, num_players=1)
 
     cell_to_state = {}
     initial_states = set()
@@ -98,11 +96,8 @@ def grid(grid: str) -> ExplicitAts:
     ats.add_states(num_states)
     ats.set_initial_states(list(initial_states))
 
-    if ats.variable_valuations is None:
-        ats.variable_valuations = EntityClassValuations()
-    if not ats.variable_valuations.has_state_valuations:
-        ats.variable_valuations.set_state_valuations(EntityValuations())
-    state_valuations = ats.variable_valuations.state_valuations
+    variable_valuations = ats.add_variable_valuations()
+    state_valuations = variable_valuations.add_state_valuations()
     vx = state_valuations.add_variable("x")
     vy = state_valuations.add_variable("y")
     for state in ats.states:

umbi/ats/explicit_ats.py

@@ -443,6 +443,12 @@ def get_observation_annotation(self) -> ObservationAnnotation:
 
     ### Variable valuations. ###
 
+    def add_variable_valuations(self) -> EntityClassValuations:
+        if self.variable_valuations is not None:
+            raise ValueError("variable valuations already exist")
+        self.variable_valuations = EntityClassValuations()
+        return self.variable_valuations
+
     @property
     def has_variable_valuations(self) -> bool:
         return self.variable_valuations is not None

umbi/ats/variable_valuations.py

@@ -315,6 +315,14 @@ def has_values_for(self, entity_class: EntityClass) -> bool:
         """Check if the annotation has values for the given entity class."""
         return self.get(entity_class) is not None
 
+    def add_valuations_for(self, entity_class: EntityClass) -> EntityValuations:
+        """Add a new EntityValuations for the given entity class."""
+        if self.has_values_for(entity_class):
+            raise ValueError(f"Annotation already has values for entity class {entity_class}")
+        valuations = EntityValuations()
+        self[entity_class] = valuations
+        return valuations
+
     def get_valuations_for(self, entity_class: EntityClass) -> EntityValuations:
         """
         Get the values for the given entity class.
@@ -382,8 +390,20 @@ def observation_valuations(self) -> EntityValuations:
     def player_valuations(self) -> EntityValuations:
         return self.get_valuations_for(EntityClass.PLAYERS)
 
-    def add_state_valuation(self):
-        return self.get_valuations_for(EntityClass.STATES)
+    def add_state_valuations(self) -> EntityValuations:
+        return self.add_valuations_for(EntityClass.STATES)
+
+    def add_choice_valuations(self) -> EntityValuations:
+        return self.add_valuations_for(EntityClass.CHOICES)
+
+    def add_branch_valuations(self) -> EntityValuations:
+        return self.add_valuations_for(EntityClass.BRANCHES)
+
+    def add_observation_valuations(self) -> EntityValuations:
+        return self.add_valuations_for(EntityClass.OBSERVATIONS)
+
+    def add_player_valuations(self) -> EntityValuations:
+        return self.add_valuations_for(EntityClass.PLAYERS)
 
     def set_state_valuations(self, values: EntityValuations):
         self.set_valuations_for(EntityClass.STATES, values)