In [ ]:
import numpy as np

# import some data
from pypge.benchmarks import diffeq

# visualization libraries
import matplotlib.pyplot as plt

# plot the visuals in ipython
%matplotlib inline

In [ ]:
# default parameters to the Simple pendulum, you may change the uncommented ones
params = {
#     'name': "SimplePendulum",
#     'xs_str': ["A", "V"],
    'params': {
        "M": 1.0,  # Mass of pendulum
        "R": 1.0   # Length of rod
    },
#     'eqn_strs': [
#         "V",               # dA
#         "(-9.8/R)*sin(A)"  # dV
#     ],
    'init_conds': {
        "A": 2.0,
        "V": 2.0
    },
    'time_end': 10.0,
    'time_step': 0.01,
    'noise': 0.1
}

# This returns the params object with more fields for Data and sympy objects
PROB = diffeq.SimplePendulum(params=params)

t_pts = PROB['time_pts']
x_pts = PROB['xs_pts']
p_pts = PROB['xs_pure']

print PROB['name']
for i,dx in enumerate(PROB['dxs_str']):
    print "  {:<4s} =  {:s}".format(dx,PROB['eqn_strs'][i])

# With Noise, Plot velocity & angle as a function of time
fig, ax1 = plt.subplots()
ax2 = ax1.twinx()
ax1.plot(t_pts, x_pts[0], 'g')
ax2.plot(t_pts, x_pts[1])

ax1.set_xlabel('time')
ax1.set_ylabel('velocity (blue)')
ax2.set_ylabel('angle (green)')

plt.show()

# No Noise, Plot velocity & angle as a function of time
fig, ax1 = plt.subplots()
ax2 = ax1.twinx()
ax1.plot(t_pts, p_pts[0], 'g')
ax2.plot(t_pts, p_pts[1])

ax1.set_xlabel('time')
ax1.set_ylabel('velocity (blue)')
ax2.set_ylabel('angle (green)')

plt.show()

In [ ]:
# Since we want a diffeq, we need to calc numerical derivatives
# [explain why eval on numerical derivative data]
dt_pts = np.gradient(t_pts, edge_order=2)
dp_pts = np.gradient(p_pts,dt_pts,edge_order=2)[1]

# But first need to smooth out the "real" data before learning
from scipy.signal import savgol_filter

win_sz = 151
poly_ord = 7
x_pts_sm = savgol_filter(x_pts, win_sz, poly_ord)

plt.plot(t_pts,x_pts[1],'b.', ms=3)
plt.plot(t_pts,x_pts_sm[1], 'r')
plt.show()

## numerical derivatives (first order)
dx_pts_sm = savgol_filter(x_pts, win_sz, poly_ord, deriv=1, delta=t_pts[1])
plt.plot(t_pts,dp_pts[1],'b.', ms=3)
plt.plot(t_pts,dx_pts_sm[1], 'r')
plt.show()

In [ ]:
# now let's do some Learning
# we will search for dV, cause that's the interesting one

from pypge.search import PGE
from pypge import expand

# create the PGE estimator
pge = PGE(
    system_type = "diffeq",
    search_vars = "y",
    usable_vars = PROB['xs_str'],
    usable_funcs = expand.BASIC_BASE[1:],
    pop_count = 3,
    peek_count = 9,
    max_iter = 4,
    workers = 2
    )

# A & V are the data values, dV is the y target
pge.fit(x_pts_sm, dx_pts_sm[1])

In [ ]:
paretos = pge.get_final_paretos()

finals = [m for front in paretos for m in front]

pge_szs = [m.size() for m in finals]
pge_scr = [m.score for m in finals]
pge_evar = [1.0 - m.evar for m in finals]

pge_szs_f = [m.size() for m in paretos[0]]
pge_scr_f = [m.score for m in paretos[0]]
pge_evar_f = [1.0 - m.evar for m in paretos[0]]

plt.plot(pge_szs, pge_scr, 'b.', pge_szs_f, pge_scr_f, 'ro')
plt.show()
plt.plot(pge_szs, pge_evar, 'b.', pge_szs_f, pge_evar_f, 'ro')
plt.show()

In [ ]:
for best_m in paretos[0]:
    print best_m
    y_pred = best_m.predict(best_m, pge.vars, x_pts)
    plt.plot(t_pts, dx_pts_sm[1], 'b.', ms=3)
    plt.plot(t_pts, y_pred, 'r')
    plt.show()