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 Chaotic pendulum
params = {
    'name': "ChaoticPendulum",
    'xs_str': ["A", "V"],
    'params': {
        "M": 1.0,   # Mass of pendulum
        "R": 1.0,   # Length of rod
        "a": 5.0,   # Amplitude of driving force
        "b": 0.05,   # Damping (friction) constant
        "k": 0.65,  # constant related to frequency of driving force
        "g": 9.81   # Gravity constant (let's got to the moon! [1.62])
    },
    'eqn_strs': [
        "V",
        "(-g/R)*sin(A) + (-b*V + a*cos(k*T) ) / (M*R**2)"
    ],
    'init_conds': {
        "A": 0.0,
        "V": 0.0
    },
    'time_end': 50.0,
    'time_step': 0.05,
    'noise': 0.05
}

# This returns the params object with more fields for Data and sympy objects
PROB = diffeq.ChaoticPendulum(**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])
ax2.plot(t_pts, x_pts[1], 'g')

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])
ax2.plot(t_pts, p_pts[1], 'g')

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 = 51
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)
win_sz = 51
poly_ord = 7
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'] +  ["T"],
    usable_funcs = expand.BASIC_BASE[1:],
    pop_count = 3,
    peek_count = 9,
    peek_npts = 200,
    max_iter = 50,
    workers = 1
    )

# Need to include time here

# (A,V,T) are the data values, dV is the y target
tt_pts = t_pts.reshape(1,len(t_pts))
print x_pts_sm.shape, tt_pts.shape
all_data = np.append(x_pts_sm, tt_pts, 0)
pge.fit(all_data[:,1:], dx_pts_sm[1,1:])

In [ ]:
#  dV   =  (-g/R)*sin(A) + (-b*V + a*cos(k*T) ) / (M*R**2)
#  looks like   a*sin(A) - b*V + c*cos(d*T)
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, all_data)
    plt.plot(t_pts, dx_pts_sm[1], 'b.', ms=3)
    plt.plot(t_pts, y_pred, 'r')
    plt.show()