From 575a33b819d29ee99fa7d4264a943043d1d64032 Mon Sep 17 00:00:00 2001
From: Daniel Friesel <daniel.friesel@uos.de>
Date: Thu, 12 Dec 2019 08:18:27 +0100
Subject: PTA.simulate: orig_state may also be a State

---
 lib/automata.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/lib/automata.py b/lib/automata.py
index 5a420fe..39ac5df 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -956,7 +956,10 @@ class PTA:
         total_duration_mae = 0.
         total_energy = 0.
         total_energy_error = 0.
-        state = self.state[orig_state]
+        if type(orig_state) is State:
+            state = orig_state
+        else:
+            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:
             if isinstance(function[0], Transition):
-- 
cgit v1.2.3