# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
import torch
import matplotlib.pyplot as plt
import copy
# %matplotlib notebook

from neurodiffeq import diff      # the differentiation operation
from neurodiffeq.ode import solve # the ANN-based solver
from neurodiffeq.conditions import IVP   # the initial condition
from neurodiffeq.networks import FCNN    # fully-connect neural network
from neurodiffeq.networks import SinActv # sin activation
from neurodiffeq.monitors import Monitor1D
from neurodiffeq.solvers import Solver1D
from neurodiffeq.callbacks import MonitorCallback
from neurodiffeq.callbacks import WeightCallback
from neurodiffeq.callbacks import WeightCallback1,WeightCallback2
from neurodiffeq.callbacks import PeriodLocal
import copy

# ### does pretraining with a kth order method work?

MAX_EPOCHS = 3000
BURNIN = 0
SUB_RATE = 1
T_EVAL = torch.linspace(-2,2,100)


# +
kth = lambda u, t: [diff(u, t) + u+u**2 +u**3 + u**4 + u**5 +1] # specify the ODE
init_val_kth = [IVP(t_0=0.0, u_0=.5)]       # specify the initial conditon
solver_kth1 = Solver1D(
    ode_system=kth,
    conditions=init_val_kth,
    t_min=-2,
    t_max=2.0,
)
solver_kth1.fit(max_epochs=MAX_EPOCHS)
pre_net = copy.deepcopy(solver_kth1.get_internals()['best_nets'])

kth = lambda u, t: [diff(u, t) - u-u**2 -u**3 - u**4 - u**5 +1] # specify the ODE
init_val_kth = [IVP(t_0=0.0, u_0=.5)]       # specify the initial conditon
solver_kth2 = Solver1D(
    ode_system=kth,
    conditions=init_val_kth,
    t_min=-2,
    t_max=2.0,
)
solver_kth2.fit(max_epochs=MAX_EPOCHS)
pre_net1 = copy.deepcopy(solver_kth2.get_internals()['best_nets'])

kth = lambda u, t: [diff(u, t) - u+u**2 -u**3 + u**4 - u**5 +1] # specify the ODE
init_val_kth = [IVP(t_0=0.0, u_0=.5)]       # specify the initial conditon
solver_kth3 = Solver1D(
    ode_system=kth,
    conditions=init_val_kth,
    t_min=-2,
    t_max=2.0,
)
solver_kth3.fit(max_epochs=MAX_EPOCHS)
pre_net2 = copy.deepcopy(solver_kth3.get_internals()['best_nets'])


# +
eps = np.arange(1,1+MAX_EPOCHS,dtype=int)[BURNIN::SUB_RATE]
MX_EPS = 100

test_sigmoid = lambda u, t: [diff(u, t) -u +u**2] # specify the ODE
init_val_sig = [IVP(t_0=0.0, u_0=.5)]     # specify the initial conditon
solver_sigmoid_test11 = Solver1D(
    ode_system=test_sigmoid,
    conditions=init_val_sig,
    t_min=-2,
    t_max=2.0,
    nets = pre_net
)
solver_sigmoid_test11.fit(max_epochs=MX_EPS)



test_sigmoid = lambda u, t: [diff(u, t) -u +u**2] # specify the ODE
init_val_sig = [IVP(t_0=0.0, u_0=.5)]     # specify the initial conditon
solver_sigmoid_test12 = Solver1D(
    ode_system=test_sigmoid,
    conditions=init_val_sig,
    t_min=-2,
    t_max=2.0,
    nets = pre_net1
)
solver_sigmoid_test12.fit(max_epochs=MX_EPS)



test_sigmoid = lambda u, t: [diff(u, t) -u +u**2] # specify the ODE
init_val_sig = [IVP(t_0=0.0, u_0=.5)]     # specify the initial conditon
solver_sigmoid_test13 = Solver1D(
    ode_system=test_sigmoid,
    conditions=init_val_sig,
    t_min=-2,
    t_max=2.0,
    nets = pre_net2
)
solver_sigmoid_test13.fit(max_epochs=MX_EPS)

test_sigmoid = lambda u, t: [diff(u, t) -u +u**2] # specify the ODE
init_val_sig = [IVP(t_0=0.0, u_0=.5)]     # specify the initial conditon
solver_sigmoid_test2 = Solver1D(
    ode_system=test_sigmoid,
    conditions=init_val_sig,
    t_min=-2,
    t_max=2.0,
)
solver_sigmoid_test2.fit(max_epochs=MX_EPS)
# -

# %matplotlib inline

plt.figure(figsize=(15,10))
ts = torch.linspace(-2,2,200).reshape(-1,1)
ss01 = solver_kth1.get_solution()(ts,to_numpy=True)
ss02 = solver_kth2.get_solution()(ts,to_numpy=True)
ss03 = solver_kth3.get_solution()(ts,to_numpy=True)
ss1 = solver_sigmoid_test11.get_solution()(ts,to_numpy=True)
ss2 = solver_sigmoid_test12.get_solution()(ts,to_numpy=True)
ss3 = solver_sigmoid_test13.get_solution()(ts,to_numpy=True)
ss_scratch = solver_sigmoid_test2.get_solution()(ts,to_numpy=True)
ss_true = torch.sigmoid(ts).detach().numpy()
ts_np = ts.detach().numpy()
plt.scatter(ts_np,ss01,label='base1',s=0.2)
plt.scatter(ts_np,ss02,label='base2',s=0.2)
plt.scatter(ts_np,ss03,label='base3',s=0.2)
plt.plot(ts_np,ss1,label='ss1')
plt.plot(ts_np,ss2,label='ss2')
plt.plot(ts_np,ss3,label='ss3')
plt.plot(ts_np,ss_scratch,label='ss_scratch')
plt.scatter(ts_np,ss_true,label='gt',s=0.2)
plt.legend()
((ss_scratch-torch.sigmoid(ts).detach().numpy())**2).mean(),((ss1-torch.sigmoid(ts).detach().numpy())**2).mean(),((ss2-torch.sigmoid(ts).detach().numpy())**2).mean(),((ss3-torch.sigmoid(ts).detach().numpy())**2).mean()



def mse(u, t):
        true_u = torch.cos(t)
        return torch.mean( (u - true_u)**2 )



equations = ['exp1','tanh1','psig1','testsig1','testsig21']
eps = [MAX_EPOCHS]


# +
param_a = np.load(f'data/{equations[-4]}/{ep}.npy')
param_b = np.load(f'data/{equations[-1]}/{ep}.npy')
# %matplotlib inline
plt.figure()
plt.hist(param_a,density=True)
plt.hist(param_b,alpha=0.3,density=True);
plt.hist(-param_b,alpha=0.3,density=True);

plt.yscale('log')
# -

meth_dex = 0
wdots= np.zeros((len(eps),len(equations)))
for meth_dex in range(len(equations)):
    for ep_index,ep in enumerate(eps):
        param_a = np.load(f'data/{equations[meth_dex]}/{ep}.npy')
        param_b = np.load(f'data/{equations[-1]}/{ep}.npy')
        wdots[ep_index,meth_dex]=np.dot(param_a,param_b)
print(wdots)


# The solution at the 100th epoch can be used to compare - but this would really be a problem with cosine and sine.. we hope that the weights learn an underlying representation useful for transfer.

((vec_exp - vec_sigtest)**2).mean()

((vec_tanh - vec_sigtest)**2).mean()

((vec_sig1 - vec_sigtest)**2).mean()

((vec_sigtest2 - vec_sigtest)**2).mean()

# MAX_EPOCHS = 3000
BURNIN = 50
equations = ['exp1','tanh1','psig1','testsig1','testsig21']
import os
# %matplotlib inline
eps = [i for i in range(BURNIN,MAX_EPOCHS) if (i%SUB_RATE==0)]

# +
meth_dex = 0
wdots= np.zeros((len(eps),len(equations)))
for meth_dex in range(len(equations)):
    for ep_index,ep in enumerate(eps):
        
        param_a = np.load(f'data/{equations[meth_dex]}/{ep}.npy')
#         print(param_a)
        param_b = np.load(f'data/{equations[-1]}/{ep}.npy')
        wdots[ep_index,meth_dex]=np.dot(param_a,param_b)/(np.linalg.norm(param_a)*np.linalg.norm(param_b))
        
plt.figure()
for i in range(len(equations)):
    plt.plot(eps,np.cumsum(np.abs(wdots),0)[:,i]/np.cumsum(np.abs(wdots),0)[:,-1],label=equations[i])
plt.legend()
# -

# ##### Ranked Weights

# +
meth_dex = 0
wdots= np.zeros((len(eps),len(equations)))
for meth_dex in range(len(equations)):
    for ep_index,ep in enumerate(eps):
        param_a = np.load(f'data/{equations[meth_dex]}/{ep}.npy')
        param_b = np.load(f'data/{equations[-1]}/{ep}.npy')
        wdots[ep_index,meth_dex]=np.dot(np.sort(param_a),np.sort(param_b))/(np.linalg.norm(param_a)*np.linalg.norm(param_b))
        
plt.figure()
for i in range(len(equations)):
    plt.plot(eps,np.cumsum(np.abs(wdots),0)[:,i]/np.cumsum(np.abs(wdots),0)[:,-1],label=equations[i])
plt.legend()
# plt.yscale('log')
# plt.xlim([2000,3000])
# plt.ylim([300,500])
# -

# using the weights or dy/dw isn't informative. we can try to use dy/dx but that is equivlent to comparing the solution. We could try a joint qp metric.



# #### 

# +
MAX_EPOCHS = 3000
BURNIN = 0
SUB_RATE = 1
T_EVAL = torch.linspace(0.5,2,1)

eps = np.arange(1,1+MAX_EPOCHS,dtype=int)[BURNIN::SUB_RATE]

mdl_weight_collector = []


cb = PeriodLocal(period=SUB_RATE)
act_cb = WeightCallback1('exp1',T_EVAL)
cb.set_action_callback(act_cb)
exponential = lambda u, t: [diff(u, t)-u+u**2] # specify the ODE
init_val_ex = [IVP(t_0=0.0, u_0=0.5)]       # specify the initial conditon
solver_exp = Solver1D(
    ode_system=exponential,
    conditions=init_val_ex,
    t_min=-2.,
    t_max=2.0,
)
solver_exp.fit(max_epochs=MAX_EPOCHS,callbacks=[cb])


# -

ts = torch.linspace(-2,2,100)
sols = solver_exp.get_solution()(ts,to_numpy=True)
plt.plot(ts.detach().numpy(),sols)

for param in (solver_exp.nets[0]).parameters():
    print(param)










StarcoderdataJupyterScriptsDedupFiltered