add automata simulation code and parameter support

2 files changed, 100 insertions, 12 deletions

diff --git a/lib/automata.py b/lib/automata.py
index 0e6cc28..c61b065 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -1,44 +1,133 @@
 class Transition:
-    def __init__(self, orig_state, dest_state, name, arguments, arg_param_map):
+    def __init__(self, orig_state, dest_state, name,
+            energy = 0, energy_function = None,
+            duration = 0, duration_function = None,
+            timeout = 0, timeout_function = None,
+            is_interrupt = False,
+            arguments = [], param_update_function = None):
         self.name = name
         self.origin = orig_state
         self.destination = dest_state
-        self.arguments = list(arguments)
+        self.energy = energy
+        self.energy_function = energy_function
+        self.duration = duration
+        self.duration_function = duration_function
+        self.timeout = timeout
+        self.timeout_function = timeout_function
+        self.is_interrupt = is_interrupt
+        self.arguments = arguments.copy()
+        self.param_update_function = param_update_function
+
+    def get_duration(self, parameters = [], args = []):
+        if self.duration_function:
+            return self.duration_function(parameters, args)
+        return self.duration
+
+    def get_energy(self, parameters = [], args = []):
+        if self.energy_function:
+            return self.energy_function(parameters, args)
+        return self.energy
+
+    def get_timeout(self, parameters = []):
+        if self.timeout_function:
+            return self.timeout_function(parameters)
+        return self.timeout
+
+    def get_params_after_transition(self, parameters, args = []):
+        if self.param_update_function:
+            return self.param_update_function(parameters, args)
+        return parameters
 
 class State:
-    def __init__(self, name):
+    def __init__(self, name, power = 0, power_function = None):
         self.name = name
-        self.outgoing_transitions = []
+        self.power = power
+        self.power_function = power_function
+        self.outgoing_transitions = {}
 
     def add_outgoing_transition(self, new_transition):
-        self.outgoing_transitions.append(new_transition)
+        self.outgoing_transitions[new_transition.name] = new_transition
+
+    def get_energy(self, duration, parameters = []):
+        if self.power_function:
+            return self.power_function(parameters) * duration
+        return self.power * duration
+
+    def get_transition(self, transition_name):
+        return self.outgoing_transitions[transition_name]
+
+    def has_interrupt_transitions(self):
+        for trans in self.outgoing_transitions.values():
+            if trans.is_interrupt:
+                return True
+        return False
+
+    def get_next_interrupt(self, parameters):
+        interrupts = filter(lambda x: x.is_interrupt, self.outgoing_transitions.values())
+        interrupts = sorted(interrupts, key = lambda x: x.get_timeout(parameters))
+        return interrupts[0]
 
     def dfs(self, depth):
         if depth == 0:
-            for trans in self.outgoing_transitions:
+            for trans in self.outgoing_transitions.values():
                 yield [trans.name]
         else:
-            for trans in self.outgoing_transitions:
+            for trans in self.outgoing_transitions.values():
                 for suffix in trans.destination.dfs(depth - 1):
                     new_suffix = [trans.name]
                     new_suffix.extend(suffix)
                     yield new_suffix
 
 class PTA:
-    def __init__(self, state_names, parameters):
+    def __init__(self, state_names = [], parameters = [], initial_param_values = None):
         self.states = dict([[state_name, State(state_name)] for state_name in state_names])
-        self.parameters = list(parameters)
+        self.parameters = parameters.copy()
+        if initial_param_values:
+            self.initial_param_values = initial_param_values.copy()
+        else:
+            self.initial_param_values = [None for x in self.parameters]
         self.transitions = []
 
         if not 'UNINITIALIZED' in state_names:
             self.states['UNINITIALIZED'] = State('UNINITIALIZED')
 
-    def add_transition(self, orig_state, dest_state, function_name, arguments, arg_param_map):
+    def add_state(self, state_name, **kwargs):
+        self.states[state_name] = State(state_name, **kwargs)
+
+    def add_transition(self, orig_state, dest_state, function_name, **kwargs):
         orig_state = self.states[orig_state]
         dest_state = self.states[dest_state]
-        new_transition = Transition(orig_state, dest_state, function_name, arguments, arg_param_map)
+        new_transition = Transition(orig_state, dest_state, function_name, **kwargs)
         self.transitions.append(new_transition)
         orig_state.add_outgoing_transition(new_transition)
 
     def dfs(self, depth = 10, orig_state = 'UNINITIALIZED'):
         return self.states[orig_state].dfs(depth)
+
+    def simulate(self, trace, orig_state = 'UNINITIALIZED'):
+        total_duration = 0.
+        total_energy = 0.
+        state = self.states[orig_state]
+        parameters = self.initial_param_values
+        for function in trace:
+            function_name = function[0]
+            function_args = function[1 : ]
+            if function_name == 'sleep':
+                duration = function_args[0]
+                total_energy += state.get_energy(duration, parameters)
+                total_duration += duration
+            else:
+                transition = state.get_transition(function_name)
+                total_duration += transition.get_duration(parameters, function_args)
+                total_energy += transition.get_energy(parameters, function_args)
+                parameters = transition.get_params_after_transition(parameters, function_args)
+                state = transition.destination
+                while (state.has_interrupt_transitions()):
+                    transition = state.get_next_interrupt(parameters)
+                    duration = transition.get_timeout(parameters)
+                    total_duration += duration
+                    total_energy += state.get_energy(duration, parameters)
+                    parameters = transition.get_params_after_transition(parameters)
+                    state = transition.destination
+
+        return total_energy, total_duration, state, parameters
diff --git a/lib/dfatool.py b/lib/dfatool.py
index 5b8f69f..6e0e829 100755
--- a/lib/dfatool.py
+++ b/lib/dfatool.py
@@ -8,7 +8,6 @@ import numpy as np
 import os
 import re
 from scipy import optimize
-from gplearn.genetic import SymbolicRegressor
 from sklearn.metrics import r2_score
 import struct
 import sys