adding kalman

kalman/plot_roukf_parameters.py

@@ -0,0 +1,121 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from itertools import cycle
+import argparse
+import pickle
+
+
+def is_ipython():
+    ''' Check if script is run in IPython.
+
+    Returns:
+        bool: True if IPython, else False '''
+    try:
+        get_ipython()
+        ipy = True
+    except NameError:
+        ipy = False
+
+    return ipy
+
+
+def load_data(file):
+    ''' Load numpy data from file.
+
+    Returns
+        dict: data dictionary
+    '''
+    dat = np.load(file)
+
+    return dat
+
+
+def plot_parameters(dat, deparameterize=False, ref=None):
+    ''' Plot the parameters in separate subplots with uncertainties.
+
+    Args:
+        dat (dict):  data dictionary
+        deparameterize (bool):  flag indicating if parameters should be
+          deparameterized via 2**theta
+        ref:    reference value to be plotted with parameters
+    '''
+    if is_ipython():
+        plt.ion()
+
+    dim = dat['theta'].shape[-1]
+    fig1, axes = plt.subplots(1, dim)
+
+    if dim == 1:
+        axes = [axes]
+
+    axes[0].set_ylabel(r'$\theta$')
+
+    t = dat['times']
+    theta = dat['theta']
+    P = dat['P_theta']
+
+    col = cycle(['C0', 'C1', 'C0', 'C1'])
+    ls = cycle(['-', '-', '--', '--', ':', ':', '-.', '-.'])
+
+    col_ = next(col)
+    ls_ = next(ls)
+    for i, ax in enumerate(axes):
+        if dim == 1:
+            theta = theta.reshape((-1, 1))
+            P = P.reshape((-1, 1, 1))
+
+        if deparameterize:
+            ax.plot(t, 2**theta[:, i], '-', c=col_, ls=ls_)
+        else:
+            ax.plot(t, theta[:, i], '-', c=col_, ls=ls_)
+            ax.fill_between(t, theta[:, i] - np.sqrt(P[:, i, i]),
+                            theta[:, i] + np.sqrt(P[:, i, i]), alpha=0.3,
+                            color=col_)
+        ax.set_xlabel(r'time')
+
+    if ref:
+        if isinstance(ref, (int, float)):
+            ref = np.array(ref)
+
+        for ax, ri in zip(axes, ref):
+            if ri:
+                # if deparameterize:
+                ax.plot((t[0], t[-1]), (ri, )*2, '-.k', lw=2,
+                        label=r'ground truth')
+                # else:
+                #     ax.plot((t[0], t[-1]), (np.log2(ri), )*2, '-.k', lw=2,
+                #             label=r'ground truth')
+
+    # print('theta_peak: \t {}'.format(theta[round(len(theta)/2), :]))
+    print('Final value theta: \t {}'.format(theta[-1, :]))
+    print('Deparameterized: 2^theta_end: \t {}'.format(2**theta[-1, :]))
+
+    if not is_ipython():
+        plt.show()
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(
+        description='''
+Plot the time evolution of the ROUKF estimated parameters.
+
+To execute in IPython::
+
+    %run plot_roukf_parameters.py [-d] [-r N [N \
+...]] file
+''',
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('file', type=str, help='path to ROUKF stats file')
+    parser.add_argument('-d', '--deparameterize', action='store_true',
+                        help='deparameterize the parameters by 2**theta')
+    parser.add_argument('-r', '--ref', metavar='N', nargs='+', default=None,
+                        type=float, help='Reference values for parameters')
+    return parser
+
+
+if __name__ == '__main__':
+    args = get_parser().parse_args()
+
+    dat = load_data(args.file)
+
+    plot_parameters(dat, deparameterize=args.deparameterize, ref=args.ref)