<jupyter_start><jupyter_text>esn_cv.Y_turbo.detach().cpu())Y_turbo = esn_cv.Y_turbo.data.cpu()plt.plot(Y_turbo)<jupyter_code>plt.plot(xtrain)
# plot_result(esn, xtrain, v0s = np.array([1]), 
#             y0s = [1.3],plot_gt = True, ode_coefs = [1,1], 
#             force_k = 0,
#            backprop_f = optimize_last_layer,solve = True, epochs = 80000, reg = False)
# plt.plot(esn.states[:,1:10].detach().cpu());
def fit_and_test(RC, xtrain, xtest, y0 = 1.3, v0 = 1, ode_coefs = [1,1], 
                 solve = None, epochs = None, reg = None, plott = None):
    train_args = {"burn_in" : int(BURN_IN), 
                      "ODE_order" : 1,
                      "force" : force,
                      "reparam_f" : reparam,
                      "init_conditions" : [float(y0), float(v0)],
                      "ode_coefs" : ode_coefs,
                      "y" : None,
                      "X" : xtrain.view(-1,1),
                      "eq_system" : True,
                      #"out_weights" : out_weights
                      }
    #fit
    y, ydot = RC.fit(**train_args, SOLVE = solve)#, out_weights = out_weights)
    states_dict = {"s"  : RC.states.clone(),
                   "s1" : RC.states_dot.clone(), 
                   "G"  : RC.G,
                   "ex" : RC.extended_states.clone(),
                   "sb1": RC.sb1,
                   "sb" : RC.sb
                   }
    weight_dict = optimize_last_layer(RC, epochs = epochs,reg = reg, plott = plott)
    RC = weight_dict["RC"]
    #y, ydot = esn.fit(**train_args, preloaded_states_dict = states_dict, out_weights = weight_dict, SOLVE = False)
    
    #test
    score, pred, _ = RC.test(y = torch.ones_like(xtest.to(esn.device)), X = xtest.to(esn.device), reparam = reparam, ODE_criterion = custom_loss)
    
    return esn.X.cpu().data, weight_dict["fit"].cpu().data, ydot.cpu().data, pred.cpu().data, weight_dict

def integrator_sol(esn):
    def f(u, t ,lam=0,A=0,W=1):
            x,  px = u      # unpack current values of u
            derivs = [px, -x - lam*x**3 +A*np.sin(W*t)]     # you write the derivative here
            return derivs
        
    # Scipy Solver   
    def NLosc_solution(t, x0,  px0, lam=0, A=0,W=1):
        u0 = [x0, px0]
        # Call the ODE solver
        solPend = odeint(f, u0, t.cpu(), args=(lam,A,W,))
        xP = solPend[:,0];        pxP = solPend[:,1];   
        return xP, pxP

    y_truth, v_truth = NLosc_solution(esn.X.squeeze().data,1.3,1,lam=1, A=0, W= 0) 
    
    return y_truth, v_truth

def plot_sol(X, yy, gt, xtest, pred, train_lim = None):
    plt.figure(figsize = (12, 5))
    print(yy[0,:].shape)
    plt.plot(X, yy[:,0].cpu(), label = "pred", color = "red")
    plt.plot(X, gt[0], '--', color = 'r')
    
    plt.axvline(train_lim, label = "train_limit")
    
    plt.plot(X, yy[:,1].cpu(), label = "pred", color = "b", linewidth = 5, alpha = 0.5)
    plt.plot(X, gt[1], '--', color = 'b', linewidth = 5, alpha = 0.5)
    
    plt.plot(xtest, pred, color = "green")
may12hps = {'dt': 0.001,
 'n_nodes': 500,
 'connectivity': 0.019946997092875757,
 'spectral_radius': 2.4289157390594482,
 'regularization': 49.04219249279563,
 'leaking_rate': 0.0032216429244726896,
 'bias': 0.3808490037918091,
 'enet_alpha': 0.2040003091096878,
 'enet_strength': 0.07488961475845243,
 'spikethreshold': 0.4231834411621094,
 'gamma': .09350859373807907,
  'gamma_cyclic' : 0.9999}

may13hps ={'gamma_cyclic': 1,#0.9998,
 'spikethreshold': 0.25,
 'enet_alpha': 0.2,
 'dt': 0.01,
 'regularization': 13.803842646028846,
 'n_nodes': 600,
 'connectivity': 0.01344268178203971,
 'spectral_radius': 2.459860324859619,
 'leaking_rate': 0.0045151556842029095,
 #'input_scaling': 0.7782557606697083,
 'bias': -0.7429814338684082,
 'enet_strength': 0.04331694643272608,
 'gamma': 0.08337975293397903}

may15hps = {'dt': 0.001,
 'regularization': 48.97788193684461,
 'n_nodes': 500,
 'connectivity': 0.017714821964432213,
 'spectral_radius': 2.3660330772399902,
 'leaking_rate': 0.0024312976747751236,
 'bias': 0.37677669525146484,
 'enet_alpha': 0.2082211971282959,
 'enet_strength': 0.118459548397668,
 'spikethreshold': 0.43705281615257263,
 'gamma': 0.09469877928495407,
 'gamma_cyclic': 0.999860422666841}

hp_set = may15hps
esn = EchoStateNetwork(**hp_set, 
                       random_state = 209, 
                       feedback = False, 
                       id_ = 10,
                       activation_f = torch.sin,
                       act_f_prime = torch.cos,
                       dtype = torch.float32, n_outputs = 2)
factor = 1#0.6*0.7
base = 10*np.pi
factor = 1.2
x0, xf, xf2 = 0, base, base*factor
nsteps = int(abs(xf - x0)/(hp_set["dt"]))
nsteps2 = int(abs(xf2 - xf)/(hp_set["dt"]))
xtrain = torch.linspace(x0, xf, nsteps, requires_grad=False).view(-1,1)
xtest = torch.cat((xtrain,xtrain+xtrain[-1]), axis = 0)[len(xtrain):]
dt1, dt2 = float(xtest[1] - xtest[0]), float(xtrain[1]- xtrain[0])
#assert dt1 == dt2, f'{dt1} != {dt2}'


xx, yy, yydot, yypred, weight_dict =  fit_and_test(esn, xtrain, xtest, epochs = 50000, reg = False, plott = True, solve = True)
plt.plot(esn.states[:,  1:10]);
xtrain[1]- xtrain[0], xtest[1]- xtest[0]
gt = integrator_sol(esn)
plot_sol(xx, yy,gt, xtest, yypred, xf)
sns.heatmap(esn.laststate.cpu().data.numpy().reshape(-1,1))
xtest.shape, weight_dict["weights"].shape
xtest2 = torch.linspace(x0, xf2, nsteps2, requires_grad=False).view(-1,1)
esn = EchoStateNetwork(**may12hps, 
                       random_state = 209, 
                       feedback = False, 
                       id_ = 10,
                       backprop = False,
                       activation_f = torch.sin,
                       act_f_prime = torch.cos,
                       dtype = torch.float32, n_outputs = 2)

train_args = {"burn_in" : int(BURN_IN), 
              "ODE_order" : 1,
              "force" : force,
              "reparam_f" : reparam,
              "init_conditions" : [float(1.3), float(1)],
              "ode_coefs" : [1,1],
              "y" : None,
              "X" : xtrain.view(-1,1),#
              "eq_system" : True,
              "out_weights" : weight_dict,
              "SOLVE" : False
              }
y, ydot = esn.fit(**train_args)
yhat = esn.predict(None, x = xtest.cuda(), continuation = True, #torch.ones_like(xtest2.cuda()), , 
                   continue_force = True)
xtrain.shape
plt.plot(esn.X.cpu(), y.cpu().detach())
plt.plot(xtest,yhat[1].cpu().detach())
esn.X_val
esn.laststate
import pickle
#t2, ys, gts, ws, bs, Ls = result
plot_data = {"time": xx, 
            "ypreds" : yy,
            "extrapolation" : yypred,
            "gts" : gt}

with open('nonlinear_oscillator_plot.pickle', 'wb') as handle:
    pickle.dump(plot_data, handle, protocol=pickle.HIGHEST_PROTOCOL)
    
with open('nonlinear_oscillator_plot.pickle', 'rb') as handle:
    b = pickle.load(handle)


repr_data = {"time": xx, 
       "hyper_params" : may12hps,
       "out_weights" : {"weights": [weight_dict["weights"]],
                        "bias": [weight_dict["bias"]]}, 
       "burn_in" : BURN_IN, 
       "epochs" : 30000,
       "learning_rate": 0.04,
       "loss_history" : weight_dict["loss"],
       "info" : "run on 30k epochs with both lr schedulers.",
       "v0" : 1,
       "y0" : 1.3}

with open('nonlinear_oscillator_reproduce.pickle', 'wb') as handle:
    pickle.dump(repr_data, handle, protocol=pickle.HIGHEST_PROTOCOL)

with open('nonlinear_oscillator_reproduce.pickle', 'rb') as handle:
    b = pickle.load(handle)
esn = EchoStateNetwork(**afternoon_hps, 
                       random_state = 209, 
                       id_ = 10,
                       activation_f = torch.sin,
                       act_f_prime = torch.cos,
                       dtype = torch.float32, n_outputs = 2)

extrapolate(esn, 0 , (np.pi * 20), (np.pi * 20)*1.2, epochs = 100000, solve = True, reg = False)
orig_BO_train_len_pi_prop = 0.6*0.7
extrapolate(esn,  0, (np.pi * 4),  (np.pi * 4)*1.2, epochs = 20000)
assert False
best_weights = {"weights" : esn.LinOut.weight.data, 
                "bias": esn.LinOut.bias.data}
xf0,xf1, dt = 0, (np.pi * 20), esn.dt
train_args = {"burn_in" : int(BURN_IN), 
              "ODE_order" : 1,
              "force" : force,
              "reparam_f" : reparam,
              "init_conditions" : [float(1.3), float(1)],
              "ode_coefs" : [1,1],
              "y" : None,
              "X" : xtrain.view(-1,1),
              "eq_system" : True,
              "out_weights" : best_weights
              }
        
y, ydot = esn.fit(**train_args, SOLVE = False)



#nsteps_test = int((xf2 - x0)/dt_)
#nsteps_test2 = int((xf2 - xf1)/dt_)

#print(f'dt = {dt_}')

#xtest = torch.linspace(x0, xf2, steps = nsteps_test, requires_grad=False).view(-1,1)
#xtest2 = torch.linspace(xf1, xf2, steps = nsteps_test2, requires_grad=False).view(-1,1)
def f(u, t ,lam=0,A=0,W=1):
            x,  px = u      # unpack current values of u
            derivs = [px, -x - lam*x**3 +A*np.sin(W*t)]     # you write the derivative here
            return derivs
        
# Scipy Solver   
def NLosc_solution(t, x0,  px0, lam=0, A=0,W=1):
    u0 = [x0, px0]
    # Call the ODE solver
    solPend = odeint(f, u0, t.cpu(), args=(lam,A,W,))
    xP = solPend[:,0];        pxP = solPend[:,1];   
    return xP, pxP

y_truth, v_truth  = NLosc_solution(esn.X.squeeze().data,1.3,1,lam=1, A=0, W= 0.5) 


#p = y[:,1].cpu()# + v0
#yy = y[:,0].cpu()# + y0
plt.plot((y[:,1].cpu()))
plt.plot(v_truth)
x, p = esn.yfit[:,0].view(-1,1), esn.yfit[:,1].view(-1,1)
xdot, pdot = esn.ydot[:,0].view(-1,1), esn.ydot[:,1].view(-1,1)
plt.plot(custom_loss(esn.X, esn.yfit, esn.ydot, None, mean = False))
plt.plot(x, label = "x")
plt.plot(p, label = "p")
plt.legend();<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16663256955810146, 'fraction_numerical': 0.07999181669394435, 'mean_word_length': 2.4018789144050103, 'pattern_counts': {'":': 7, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '497752', 'license': 'MIT', 'n_tokens_mistral': 4346, 'n_tokens_neox': 3842, 'n_words': 717, 'path': 'final_notebooks/.ipynb_checkpoints/final_systems_BO-checkpoint.ipynb', 'repo_name': 'blindedjoy/RcTorch', 'size': '9776'}