PTA: Use "None" transition for sleep

This way, PTA models can have sleep transitions without conflict

2 files changed, 19 insertions, 5 deletions

diff --git a/lib/automata.py b/lib/automata.py
index 50ae179..1ad80bb 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -624,14 +624,28 @@ class PTA:
             return generator
 
     def simulate(self, trace: list, orig_state: str = 'UNINITIALIZED'):
+        """
+        Simulate a single run through the PTA and return total energy, duration, final state, and resulting parameters.
+
+        :param trace: list of (function name, arg1, arg2, ...) tuples representing the individual transitions,
+            or list of (Transition, argument tuple, parameter) tuples originating from dfs.
+            The tuple (None, duration) represents a sleep time between states in us
+        :param orig_state: origin state, default UNINITIALIZED
+
+        :returns (total energy in uJ, total duration in us, end state, end parameters)
+        """
         total_duration = 0.
         total_energy = 0.
         state = self.state[orig_state]
         param_dict = dict([[self.parameters[i], self.initial_param_values[i]] for i in range(len(self.parameters))])
         for function in trace:
-            function_name = function[0]
-            function_args = function[1 : ]
-            if function_name == 'sleep':
+            if isinstance(function[0], Transition):
+                function_name = function[0].name
+                function_args = function[1]
+            else:
+                function_name = function[0]
+                function_args = function[1 : ]
+            if function_name is None:
                 duration = function_args[0]
                 total_energy += state.get_energy(duration, param_dict)
                 total_duration += duration
diff --git a/test/test_pta.py b/test/test_pta.py
index dba67d7..e8de399 100755
--- a/test/test_pta.py
+++ b/test/test_pta.py
@@ -328,9 +328,9 @@ class TestPTA(unittest.TestCase):
         pta.add_transition('TX', 'IDLE', 'txComplete', timeout = 2000, is_interrupt = True)
         trace = [
             ['init'],
-            ['sleep', 10000000],
+            [None, 10000000],
             ['send', 'foo', 3],
-            ['sleep', 5000000],
+            [None, 5000000],
             ['send', 'foo', 3]
         ]
         expected_energy = 5. * 10000000 + 3 + 100 * 2000 + 5 * 5000000 + 3 + 100 * 2000