# ---
# 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
from scipy.integrate import RK45,solve_ivp
from ODE_potentials import VanDerPolePotential,LotkiVolterraPotential
from ODE_samplers import MALA_ODE,ULA_ODE,grad_ascent_ODE,run_eval_test,set_function
from multiprocessing import Pool
import multiprocessing
from zv_cv import Eval_ZVCV
from sklearn.preprocessing import PolynomialFeatures
import copy

from baselines import construct_ESVM_kernel, construct_Tukey_Hanning 
from optimize import optimize_parallel_new
from utils import *
# -

# Parameters for van-der-Pole and Lotka-Volterra examples:

# +
typ = 'LV' #'LV' for Lotka-Volterra, 'VdP' for Van-der-Pole
method = {"sampler":"MALA"} #switch between ULA and MALA
f_type = "sum_comps"

if typ == 'VdP':
    #true parameter value
    theta_star = 1.0
    #initial coordiante and speed
    y0 = np.array([0.0,2.0],dtype=float)
    #error of measurements
    sigma = 0.5
    #prior variance
    sigma_prior = 0.5
elif typ == 'LV':
    theta_star = np.array([0.6,0.025,0.8,0.025],dtype = float)
    #initial number of victims and predators
    y0 = np.array([30.0,4.0],dtype=float)
    #setting prior parameters
    sigma_prior = np.array([0.5,0.05,0.5,0.05],dtype = float)
    mu_prior = np.array([1.0,0.05,1.0,0.05],dtype=float)
    #measurements error
    sigma = np.array([0.25,0.25])
# -

# Timestaps

#initial and last time moments
t0 = 0
t_bound = 10
N_steps = 10
#moments of observations
t_moments = np.linspace(t0,t_bound,N_steps+1)
print(t_moments)

# Creating potentials

if typ == 'VdP':
    Cur_pot = VanDerPolePotential(sigma,sigma_prior,t_moments,theta_star,y0,t0,t_bound)
elif typ == 'LV':
    Cur_pot = LotkiVolterraPotential(sigma,mu_prior,sigma_prior,t_moments,theta_star,y0,t0,t_bound)

# Sampling (currently with MALA)

r_seed = 666
#burn-in period
N_burn = 1*10**3
#Train size
N_train = 1*10**4
#Test size
N_test = 1*10**4
#number of test trajectories
n_traj = 1
if typ == 'VdP':
    #dimension
    d = 1
    #step size
    step = 1e-3
elif typ == 'LV':
    #dimension
    d = 4
    #step size
    step = 5e-6

# ### Construct kernels and sample

if typ == 'VdP':
    params_prior = {"sigma":sigma_prior}
elif typ == 'LV':
    params_prior = {"mu":mu_prior,"sigma":sigma_prior}

# ### Compute starting point (maximum likelihood)

N_steps_ascent = 5000
traj,traj_grad = grad_ascent_ODE(1453,Cur_pot,step,params_prior,N_steps_ascent,d,typ)
theta_mle = traj[-1,:]
print("mle for parameters: ",theta_mle)
Cur_pot.set_mle(theta_mle)

# ### Setting function 

#step_train = 5e-7
inds_arr = np.array([0])
params = {"ind":0}
t_moments = None
r_seed = 777
traj = []
traj_grad = []
#generate data
nbcores = multiprocessing.cpu_count()
trav = Pool(nbcores)
res = trav.starmap(ULA_ODE, [(r_seed+i,Cur_pot, step, params_prior, N_burn, N_train, d, typ) for i in range (n_traj)])
trav.close()
for i in range(len(res)):
    traj.append(res[i][0])
    traj_grad.append(res[i][1])
    #print("accepted = ",res[i][2])
traj = np.asarray(traj)
traj_grad = np.asarray(traj_grad)
traj_grad = (-1)*traj_grad


# ### Apply control variates

# +
def H(k, x):
    if k==0:
        return 1.0
    if k ==1:
        return x
    if k==2:
        return (x**2 - 1)/np.sqrt(2)
    c = np.zeros(k+1,dtype = float)
    c[k] = 1.0
    h = P.hermite_e.hermeval(x,c) / np.sqrt(sp.special.factorial(k)) 
    return h

def compute_H(k,x,d):
    cur_prod = 1.0
    for i in range(d):
        cur_prod = cur_prod*H(k[i],x[:,i])
    return cur_prod


# +
def generate_lexicographical_2nd_order(d):
    """
    function to generate lexigoraphical polynomials of 1st and 2nd order
    """
    cur_list = []
    for i in range(d):
        cur_permute = np.zeros(d,dtype = int)
        cur_permute[i] = 1
        cur_list.append(np.copy(cur_permute))
        cur_permute[i] = 0
    for i in range(d):
        cur_permute = np.zeros(d,dtype = int)
        cur_permute[i] = 2
        cur_list.append(np.copy(cur_permute))
        cur_permute[i] = 1
        for j in range(i+1,d):
            cur_permute[j] = 1
            cur_list.append(np.copy(cur_permute))
            cur_permute[j] = 0
    return cur_list

a = generate_lexicographical_2nd_order(3)
print(a)


# -

def test_traj(coefs_poly_regr,gamma,r_seed,lag,d,N_test,x0):
    """
    function to perform 1-dimensional martingale decomposition
    """
    X_test,Noise = generate_traj(x0,N_test,gamma,r_seed,d)
    test_stat_vanilla = np.zeros(N_test,dtype = float)
    test_stat_vr = np.zeros_like(test_stat_vanilla)
    #compute number of basis polynomials
    basis_funcs = generate_lexicographical_2nd_order(d)
    num_basis_funcs = len(basis_funcs)
    #compute polynomials of noise variables Z_l
    poly_vals = np.zeros((num_basis_funcs,N_test), dtype = float)
    for k in range(len(basis_funcs)):
        poly_vals[k,:] = compute_H(basis_funcs[k],Noise,d)
    #initialize function
    f_vals_vanilla = np.sum(X_test**2,axis=1)
    #array to store control variates values
    cvfs = np.zeros_like(f_vals_vanilla)
    #compute coeffitients bar_a
    bar_a_1st_order = np.zeros((lag,d,N_test),dtype=float)
    bar_a_2nd_order = np.zeros((lag,d,d,N_test),dtype=float)
    #preprocessing 
    X_test = np.concatenate((x0.reshape(1,d),X_test),axis=0)
    coefs_poly_1st_order = np.zeros((lag,d),dtype=float)
    coefs_poly_1st_order = coefs_poly_regr[:,1:d+1]
    coefs_poly_2nd_order = np.zeros((lag,d,d),dtype=float)
    counter = 0
    for i in range(d):
        for j in range(i,d):
            coefs_poly_2nd_order[:,i,j] = coefs_poly_regr[:,d+1+counter]
            counter += 1
    for i in range(lag):
        #first-order coefficients
        for j in range(d):
            bar_a_1st_order[i,j,:] = coefs_poly_1st_order[i,j]*np.sqrt(gamma)*sigma(X_test[:-1])[:,j]
            #sum more coefficients
            for k in range(j):
                bar_a_1st_order[i,j,:] += coefs_poly_2nd_order[i,k,j]*np.sqrt(gamma)*sigma(X_test[:-1])[:,j]*((X_test[:-1]+gamma*b(X_test[:-1]))[:,k])
            #diagonal part
            bar_a_1st_order[i,j] += 2*coefs_poly_2nd_order[i,j,j]*np.sqrt(gamma)*sigma(X_test[:-1])[:,j]*((X_test[:-1]+gamma*b(X_test[:-1]))[:,j])
            #sum more coefficients
            for k in range(j+1,d):
                bar_a_1st_order[i,j] += coefs_poly_2nd_order[i,j,k]*np.sqrt(gamma)*sigma(X_test[:-1])[:,j]*((X_test[:-1]+gamma*b(X_test[:-1]))[:,k])
            #second-order coefficients, to be filled
            bar_a_2nd_order[i,j,j,:] = coefs_poly_2nd_order[i,j,j]*np.sqrt(2)*gamma*sigma(X_test[:-1])[:,j]
            for k in range(j+1,d):
                bar_a_2nd_order[i,j,k,:] = coefs_poly_2nd_order[i,j,k]*gamma*sigma(X_test[:-1])[:,j]*sigma(X_test[:-1])[:,k]
        """
        bar_a_0_1[i,1:] = coefs_poly_regr[i,1]*cov[0,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,0]+\
                        coefs_poly_regr[i,2]*cov[1,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,1]+\
                        2*coefs_poly_regr[i,3]*cov[0,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,0]*(X_test[:-1]+gamma*b(X_test[:-1]))[:,0]+\
                        coefs_poly_regr[i,4]*(((X_test[:-1]+gamma*b(X_test[:-1]))[:,0])*sigma(X_test[:-1])[:,1]*np.sqrt(gamma)*cov[1,1] +\
                                             ((X_test[:-1]+gamma*b(X_test[:-1]))[:,1])*sigma(X_test[:-1])[:,0]*np.sqrt(gamma)*cov[0,1])+\
                        2*coefs_poly_regr[i,5]*cov[1,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,1]*(X_test[:-1]+gamma*b(X_test[:-1]))[:,1]
        bar_a_0_1[i,0] = coefs_poly_regr[i,1]*cov[0,1]*np.sqrt(gamma)*sigma(x0)[0]+\
                        coefs_poly_regr[i,2]*cov[1,1]*np.sqrt(gamma)*sigma(x0)[1]+\
                        2*coefs_poly_regr[i,3]*cov[0,1]*np.sqrt(gamma)*sigma(x0)[0]*(x0+gamma*b(x0))[0]+\
                        coefs_poly_regr[i,4]*(((x0+gamma*b(x0))[0])*sigma(x0)[1]*np.sqrt(gamma)*cov[1,1] +\
                                             ((x0+gamma*b(x0))[1])*sigma(x0)[0]*np.sqrt(gamma)*cov[0,1])+\
                        2*coefs_poly_regr[i,5]*cov[1,1]*np.sqrt(gamma)*sigma(x0)[1]*(x0+gamma*b(x0))[1]
        #coefficients with H_1_0
        bar_a_1_0[i,1:] = coefs_poly_regr[i,1]*cov[0,0]*np.sqrt(gamma)*sigma(X_test[:-1])[:,0]+\
                        coefs_poly_regr[i,2]*cov[0,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,1]+\
                        2*coefs_poly_regr[i,3]*cov[0,0]*np.sqrt(gamma)*sigma(X_test[:-1])[:,0]*(X_test[:-1]+gamma*b(X_test[:-1]))[:,0]+\
                        coefs_poly_regr[i,4]*(((X_test[:-1]+gamma*b(X_test[:-1]))[:,0])*sigma(X_test[:-1])[:,1]*np.sqrt(gamma)*cov[0,1] +\
                                             ((X_test[:-1]+gamma*b(X_test[:-1]))[:,1])*sigma(X_test[:-1])[:,0]*np.sqrt(gamma)*cov[0,0])+\
                        2*coefs_poly_regr[i,5]*cov[0,1]*np.sqrt(gamma)*sigma(X_test[:-1])[:,1]*(X_test[:-1]+gamma*b(X_test[:-1]))[:,1]
        bar_a_1_0[i,0] = coefs_poly_regr[i,1]*cov[0,0]*np.sqrt(gamma)*sigma(x0)[0]+\
                        coefs_poly_regr[i,2]*cov[0,1]*np.sqrt(gamma)*sigma(x0)[1]+\
                        2*coefs_poly_regr[i,3]*cov[0,0]*np.sqrt(gamma)*sigma(x0)[0]*(x0+gamma*b(x0))[0]+\
                        coefs_poly_regr[i,4]*(((x0+gamma*b(x0))[0])*sigma(x0)[1]*np.sqrt(gamma)*cov[0,1] +\
                                             ((x0+gamma*b(x0))[1])*sigma(x0)[0]*np.sqrt(gamma)*cov[0,0]) +\
                        2*coefs_poly_regr[i,5]*cov[0,1]*np.sqrt(gamma)*sigma(x0)[1]*(x0+gamma*b(x0))[1]
        #second-order coefficients
        bar_a_1_1[i,1:] = 2*coefs_poly_regr[i,3]*gamma*(sigma(X_test[:-1])[:,0])**2*cov[0,0]*cov[0,1]+\
                        coefs_poly_regr[i,4]*gamma*(sigma(X_test[:-1])[:,0])*(sigma(X_test[:-1])[:,1])*(cov[0,1]**2 + cov[0,0]*cov[1,1])+\
                        2*coefs_poly_regr[i,5]*gamma*(sigma(X_test[:-1])[:,1])**2*cov[1,1]*cov[0,1]
        bar_a_1_1[i,0] = 2*coefs_poly_regr[i,3]*gamma*(sigma(x0)[0])**2*cov[0,0]*cov[0,1]+\
                        coefs_poly_regr[i,4]*gamma*(sigma(x0)[0])*(sigma(x0)[1])*(cov[0,1]**2 + cov[0,0]*cov[1,1])+\
                        2*coefs_poly_regr[i,5]*gamma*(sigma(x0)[1])**2*cov[1,1]*cov[0,1]
        #coefficients with H_2_0
        bar_a_2_0[i,1:] = np.sqrt(2)*coefs_poly_regr[i,3]*gamma
        #coefficients with H_0_2
        bar_a_0_2[i,1:] = np.sqrt(2)*coefs_poly_regr[i,5]*gamma
        """
    #bar_a_1_0 = bar_a_1_0*poly_vals[0,:]
    #bar_a_0_1 = bar_a_0_1*poly_vals[1,:]
    #bar_a_1_1 = bar_a_1_1*poly_vals[2,:]
    #bar_a_2_0 = bar_a_2_0*poly_vals[3,:]
    #bar_a_0_2 = bar_a_0_2*poly_vals[4,:]
    counter = 0
    for i in range(d):
        bar_a_1st_order[:,i,:] = poly_vals[i,:]*bar_a_1st_order[:,i,:]
        bar_a_2nd_order[:,i,i,:] = poly_vals[d + d*(d+1)//2 - (d-i)*(d-i+1)//2,:]*bar_a_2nd_order[:,i,i,:]
        counter += 1
        for j in range(i+1,d):
            bar_a_2nd_order[:,i,j,:] = poly_vals[d+counter]*bar_a_2nd_order[:,i,j,:]
            counter += 1
    #compute martingale sums
    M_n_1st = 0.0
    M_n_2nd = 0.0
    for l in range(N_test):
        for r in range(min(N_test-l,lag)):
            M_n_1st += np.sum(bar_a_1st_order[r,:,l])
            M_n_2nd += np.sum(bar_a_2nd_order[r,:,:,l])
    return np.mean(f_vals_vanilla), np.mean(f_vals_vanilla)-(M_n_1st)/N_test, np.mean(f_vals_vanilla)-(M_n_1st+M_n_2nd)/N_test
#
#
#


def approx_q(X_train,Y_train,N_traj_train,lag,max_deg):
    """
    Function to regress q functions on a polynomial basis;
    Args:
        X_train - train tralectory;
        Y_train - function values;
        N_traj_train - number of training trajectories;
        lag - truncation point for coefficients, those for |p-l| > lag are set to 0;
        max_deg - maximum degree of polynomial in regression
    """
    dim = X_train[0,:].shape[0]
    #print("dimension = ",dim)
    coefs_poly = np.array([])
    for i in range(lag):
        x_all = np.array([])
        y_all = np.array([])
        for j in range(N_traj_train):
            y = Y_train[j,i:,0]
            if i == 0:
                x = X_train[j,:]
            else:
                x = X_train[j,:-i]
            #concatenate results
            if x_all.size == 0:
                x_all = x
            else:
                x_all = np.concatenate((x_all,x),axis = 0)
            y_all = np.concatenate([y_all,y])
        #should use polyfeatures here
        #print("variance: ",np.var(y_all))
        #print(y_all[:50])
        poly = PolynomialFeatures(max_deg)
        X_features = poly.fit_transform(x_all)
        #print(X_features.shape)
        lstsq_results = np.linalg.lstsq(X_features,y_all,rcond = None)
        coefs = copy.deepcopy(lstsq_results[0])
        coefs.resize((1,X_features.shape[1]))           
        if coefs_poly.size == 0:
            coefs_poly = copy.deepcopy(coefs)
        else:
            coefs_poly = np.concatenate((coefs_poly,coefs),axis=0)
    return coefs_poly


X_train = traj
print(X_train.shape)
Y_train = X_train[:,:,0].reshape((1,-1,1))
print(Y_train.shape)
lag = 100
S_max = 2

#polynomial coefficients
coefs_poly = approx_q(X_train,Y_train,n_traj,lag,S_max)
#print(coefs_poly.shape)
print(coefs_poly)
regr_vals = np.zeros((lag,X_train.shape[1]),dtype=float)
poly = PolynomialFeatures(S_max)
features = poly.fit_transform(X_train[0])
#features = np.zeros((X_train.shape[1],6),dtype=float)
#features[:,0] = 1.0
#features[:,1:3] = X_train[0,:,:]
#features[:,3] = X_train[0,:,0]**2 
#features[:,4] = X_train[0,:,0]*X_train[0,:,1]
#features[:,5] = X_train[0,:,1]**2
for i in range(len(regr_vals)):
    regr_vals[i,:] = np.sum(coefs_poly[i,:]*features,axis=1)

# Test our regressors

cur_lag = 99
N_pts = 500
plt.figure(figsize=(10, 10))
plt.title("Testing regression model",fontsize=20)
plt.plot(Y_train[0,cur_lag:N_pts+cur_lag,0],color='r',label='true function')
plt.plot(regr_vals[cur_lag,:N_pts],color='g',label = 'practical approximation')
plt.legend(loc = 'upper left',fontsize = 16)
plt.show()



StarcoderdataJupyterScriptsDedupFiltered