support plotting of raw and parameterized data + custom functions

2 files changed, 133 insertions, 2 deletions

diff --git a/lib/dfatool.py b/lib/dfatool.py
index a03abc4..a591be3 100755
--- a/lib/dfatool.py
+++ b/lib/dfatool.py
@@ -50,6 +50,45 @@ def vprint(verbose, string):
     if verbose:
         print(string)
 
+    def _gplearn_add_(x, y):
+        return x + y
+
+    def _gplearn_sub_(x, y):
+        return x - y
+
+    def _gplearn_mul_(x, y):
+        return x * y
+
+    def _gplearn_div_(x, y):
+        if np.abs(y) > 0.001:
+            return x / y
+        return 1.
+
+def gplearn_to_function(function_str):
+
+    eval_globals = {
+        'add' : lambda x, y : x + y,
+        'sub' : lambda x, y : x - y,
+        'mul' : lambda x, y : x * y,
+        'div' : lambda x, y : np.divide(x, y) if np.abs(y) > 0.001 else 1.,
+        'sqrt': lambda x : np.sqrt(np.abs(x)),
+        'log' : lambda x : np.log(np.abs(x)) if np.abs(x) > 0.001 else 0.,
+        'inv' : lambda x : 1. / x if np.abs(x) > 0.001 else 0.,
+    }
+
+    last_arg_index = 0
+    for i in range(0, 100):
+        if function_str.find('X{:d}'.format(i)) >= 0:
+            last_arg_index = i
+
+    arg_list = []
+    for i in range(0, last_arg_index+1):
+        arg_list.append('X{:d}'.format(i))
+
+    eval_str = 'lambda {}, *whatever: {}'.format(','.join(arg_list), function_str)
+    print(eval_str)
+    return eval(eval_str, eval_globals)
+
 def _elem_param_and_arg_list(elem):
     param_dict = elem['parameter']
     paramkeys = sorted(param_dict.keys())
@@ -446,7 +485,7 @@ class AnalyticFunction:
         self._function = eval('lambda reg_param, model_param: ' + rawfunction);
         self._regression_args = list(np.ones((num_vars)))
 
-    def _get_fit_data(self, by_param, state_or_tran, model_attribute):
+    def get_fit_data(self, by_param, state_or_tran, model_attribute):
         dimension = len(self._parameter_names) + self._num_args
         X = [[] for i in range(dimension)]
         Y = []
@@ -477,7 +516,7 @@ class AnalyticFunction:
         return X, Y, num_valid, num_total
 
     def fit(self, by_param, state_or_tran, model_attribute):
-        X, Y, num_valid, num_total = self._get_fit_data(by_param, state_or_tran, model_attribute)
+        X, Y, num_valid, num_total = self.get_fit_data(by_param, state_or_tran, model_attribute)
         if num_valid > 2:
             error_function = lambda P, X, y: self._function(P, X) - y
             try:
@@ -750,6 +789,7 @@ class EnergyModel:
         self.by_param = {}
         self.by_trace = {}
         self.stats = {}
+        self.cache = {}
         np.seterr('raise')
         self._parameter_names = sorted(self.traces[0]['trace'][0]['parameter'].keys())
         self._num_args = {}
@@ -768,6 +808,11 @@ class EnergyModel:
         self._aggregate_to_ndarray(self.by_name)
         self._compute_all_param_statistics()
 
+    def distinct_param_values(self, state_or_tran, param_index = None, arg_index = None):
+        if param_index != None:
+            param_values = map(lambda x: x[param_index], self.by_name[state_or_tran]['param'])
+        return sorted(set(param_values))
+
     def _compute_outlier_stats(self, ignore_trace_indexes, threshold):
         tmp_by_param = {}
         self.median_by_param = {}
@@ -934,7 +979,16 @@ class EnergyModel:
     def get_param_analytic(self):
         static_model = self._get_model_from_dict(self.by_name, np.median)
 
+    def param_index(self, param_name):
+        if param_name in self._parameter_names:
+            return self._parameter_names.index(param_name)
+        return len(self._parameter_names) + int(param_name)
+
     def get_fitted(self):
+
+        if 'fitted_model_getter' in self.cache and 'fitted_info_getter' in self.cache:
+            return self.cache['fitted_model_getter'], self.cache['fitted_info_getter']
+
         static_model = self._get_model_from_dict(self.by_name, np.median)
         param_model = dict([[state_or_tran, {}] for state_or_tran in self.by_name.keys()])
         fit_queue = []
@@ -1018,6 +1072,9 @@ class EnergyModel:
                 return param_model[name][key]
             return None
 
+        self.cache['fitted_model_getter'] = model_getter
+        self.cache['fitted_info_getter'] = info_getter
+
         return model_getter, info_getter
 
 
diff --git a/lib/plotter.py b/lib/plotter.py
index 784ba56..6ee2691 100755
--- a/lib/plotter.py
+++ b/lib/plotter.py
@@ -5,6 +5,14 @@ import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.patches import Polygon
 
+def float_or_nan(n):
+    if n == None:
+        return np.nan
+    try:
+        return float(n)
+    except ValueError:
+        return np.nan
+
 def flatten(somelist):
     return [item for sublist in somelist for item in sublist]
 
@@ -71,6 +79,72 @@ def plot_substate_thresholds_p(model, aggregate):
     data = [aggregate[key]['sub_thresholds'] for key in keys]
     boxplot(keys, None, None, data, 'Zustand', '% Clipping')
 
+def plot_y(Y, ylabel = None, title = None):
+    plot_xy(np.arange(len(Y)), Y, ylabel = ylabel, title = title)
+
+def plot_xy(X, Y, xlabel = None, ylabel = None, title = None):
+    fig, ax1 = plt.subplots(figsize=(10,6))
+    if title != None:
+        fig.canvas.set_window_title(title)
+    if xlabel != None:
+        ax1.set_xlabel(xlabel)
+    if ylabel != None:
+        ax1.set_ylabel(ylabel)
+    plt.subplots_adjust(left = 0.05, bottom = 0.05, right = 0.99, top = 0.99)
+    plt.plot(X, Y, "rx")
+    plt.show()
+
+def _param_slice_eq(a, b, index):
+    return (*a[1][:index], *a[1][index+1:]) == (*b[1][:index], *b[1][index+1:]) and a[0] == b[0]
+
+def plot_param(model, state_or_trans, attribute, param_idx, xlabel = None, ylabel = None, title = None, extra_functions = []):
+    fig, ax1 = plt.subplots(figsize=(10,6))
+    if title != None:
+        fig.canvas.set_window_title(title)
+    if xlabel != None:
+        ax1.set_xlabel(xlabel)
+    if ylabel != None:
+        ax1.set_ylabel(ylabel)
+    plt.subplots_adjust(left = 0.05, bottom = 0.05, right = 0.99, top = 0.99)
+
+    param_model, param_info = model.get_fitted()
+
+    by_other_param = {}
+
+    for k, v in model.by_param.items():
+        if k[0] == state_or_trans:
+            other_param_key = (*k[1][:param_idx], *k[1][param_idx+1:])
+            if not other_param_key in by_other_param:
+                by_other_param[other_param_key] = {'X': [], 'Y': []}
+            by_other_param[other_param_key]['X'].extend([float(k[1][param_idx])] * len(v[attribute]))
+            by_other_param[other_param_key]['Y'].extend(v[attribute])
+
+    cm = plt.get_cmap('brg', len(by_other_param))
+    for i, k in enumerate(by_other_param):
+        v = by_other_param[k]
+        v['X'] = np.array(v['X'])
+        v['Y'] = np.array(v['Y'])
+        plt.plot(v['X'], v['Y'], "rx", color=cm(i))
+        x_range = int((v['X'].max() - v['X'].min()) * 2)
+        xsp = np.linspace(v['X'].min(), v['X'].max(), x_range)
+        if param_model:
+            ysp = []
+            for x in xsp:
+                xarg = [*k[:param_idx], x, *k[param_idx:]]
+                ysp.append(param_model(state_or_trans, attribute, param = xarg))
+            plt.plot(xsp, ysp, "r-", color=cm(i), linewidth=0.5)
+        if len(extra_functions) != 0:
+            for f in extra_functions:
+                ysp = []
+                with np.errstate(divide='ignore', invalid='ignore'):
+                    for x in xsp:
+                        xarg = [*k[:param_idx], x, *k[param_idx:]]
+                        ysp.append(f(*xarg))
+                plt.plot(xsp, ysp, "r--", color=cm(i), linewidth=1, dashes=(3, 3))
+
+    plt.show()
+
+
 def plot_param_fit(function, name, fitfunc, funp, parameters, datatype, index, X, Y, xaxis=None, yaxis=None):
     fig, ax1 = plt.subplots(figsize=(10,6))
     fig.canvas.set_window_title("fit %s" % (function))