move PTA simulation result to separate class

2 files changed, 50 insertions, 41 deletions

diff --git a/lib/automata.py b/lib/automata.py
index 6ace61f..eb52081 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -13,6 +13,15 @@ def _dict_to_list(input_dict: dict) -> list:
     return [input_dict[x] for x in sorted(input_dict.keys())]
 
 
+class SimulationResult:
+    def __init__(self, duration: float, energy: float, end_state, parameters):
+        self.duration = duration * 1e-6
+        self.energy = energy * 1e-12
+        self.end_state = end_state
+        self.parameters = parameters
+        self.mean_power = self.energy / self.duration
+
+
 class PTAAttribute:
     def __init__(self, value: float = 0, function: AnalyticFunction = None, value_error=None, function_error=None):
         self.value = value
@@ -957,7 +966,7 @@ class PTA:
                     param_dict = transition.get_params_after_transition(param_dict)
                     state = transition.destination
 
-        return total_energy, total_duration, state, param_dict
+        return SimulationResult(total_duration, total_energy, state, param_dict)
 
     def update(self, static_model, param_model, static_error=None, analytic_error=None):
         for state in self.state.values():
diff --git a/test/test_pta.py b/test/test_pta.py
index 2eb7139..cf8bce4 100755
--- a/test/test_pta.py
+++ b/test/test_pta.py
@@ -571,11 +571,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 5. * 10000000 + 3 + 100 * 2000 + 5 * 5000000 + 3 + 100 * 2000
         expected_duration = 50000 + 10000000 + 10 + 2000 + 5000000 + 10 + 2000
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {})
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {})
 
     def test_simulation_param_none(self):
         pta = PTA(parameters=['txpower', 'length'])
@@ -589,11 +589,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 500000
         expected_duration = 50000
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': None,
             'length': None
         })
@@ -613,11 +613,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 510000
         expected_duration = 50120
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': None
         })
@@ -637,11 +637,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 510000
         expected_duration = 50120
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': None
         })
@@ -656,11 +656,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 500000
         expected_duration = 50000
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': None
         })
@@ -683,11 +683,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 500000 + 10000 + (3 + 5 * 3) + (2000 * 100)
         expected_duration = 50000 + 120 + (48 + 8 * 3) + 2000
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': None
         })
@@ -709,11 +709,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 500000 + 10000 + (3 + 5 * 6) + (2000 * 100)
         expected_duration = 50000 + 120 + (48 + 8 * 6) + 2000
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': None
         })
@@ -738,11 +738,11 @@ class TestPTA(unittest.TestCase):
         ]
         expected_energy = 500000 + 10000 + (3 + 5 * 3) + (1000 + 2 * 10) * (500 + 16 * 3)
         expected_duration = 50000 + 120 + 10 + (500 + 16 * 3)
-        energy, duration, state, parameters = pta.simulate(trace)
-        self.assertEqual(energy, expected_energy)
-        self.assertEqual(duration, expected_duration)
-        self.assertEqual(state.name, 'IDLE')
-        self.assertEqual(parameters, {
+        result = pta.simulate(trace)
+        self.assertAlmostEqual(result.energy, expected_energy * 1e-12, places=12)
+        self.assertAlmostEqual(result.duration, expected_duration * 1e-6, places=6)
+        self.assertEqual(result.end_state.name, 'IDLE')
+        self.assertEqual(result.parameters, {
             'txpower': 10,
             'length': 3
         })