# ---
# 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
# ---

# +
#data format library
import h5py
#numpy
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
# # %matplotlib notebook
from sklearn.cluster import KMeans

import sys
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.colors as colors
import os
from scipy.integrate import odeint
sys.path.append('../utils/')
import operator_calculations as op_calc
import stats
from sklearn.linear_model import LinearRegression
import delay_embedding as embed
import clustering_methods as cl
colors_state_=plt.rcParams['axes.prop_cycle'].by_key()['color'][:10]
plt.rc('text', usetex=True)
plt.rc('font',size=14)


# -

def Lorenz(state,t,sigma,rho,beta):
    # unpack the state vector
    x,y,z = state
    # compute state derivatives
    xd = sigma * (y-x)
    yd = (rho-z)*x - y
    zd = x*y - beta*z
    # return the state derivatives
    return [xd, yd, zd]
dt = 0.02
frameRate=1/dt
T = 5000
state0 = np.array([-8, -8, 27])
t = np.linspace(0, T, int(T*frameRate))
sigma,rho,beta=10,28,8/3
tseries=np.array(odeint(Lorenz,state0,t,args=(sigma,rho,beta)),dtype=np.float64)[int(len(t)/2):]

plt.plot(tseries[:,0],tseries[:,2],lw=.02)
plt.show()

# # Compute predictability as a function of delay

X = tseries[:,0].reshape((tseries.shape[0],1)) #take x variable only

# +
#to get error estimates in the manuscript we split the trajectory into non-overlapping segments

n_seed_range=np.arange(200,1100,200) #number of partitions to examine
range_Ks =  np.arange(1,12,dtype=int) #range of delays to study
h_K=np.zeros((len(range_Ks),len(n_seed_range)))
for k,K in enumerate(range_Ks):
    traj_matrix = embed.trajectory_matrix(X,K=K-1)
    for ks,n_seeds in enumerate(n_seed_range):
        labels=cl.kmeans_knn_partition(traj_matrix,n_seeds)
        h = op_calc.get_entropy(labels)
        h_K[k,ks]=h
        print('Computed for {} delays and {} seeds.'.format(K,n_seeds))
# -

plt.plot(n_seed_range,h_K.T)
plt.xlabel('N',fontsize=15)
plt.ylabel('h (nats/s)',fontsize=15)
plt.show()

plt.plot(h_K[:,-1],marker='o')
plt.show()

K_star=6
traj_matrix = embed.trajectory_matrix(X,K=K_star-1)

# # Estimate implied time scales of the reversibilized operator

# +
#to get error estimates in the manuscript we split the trajectory into non-overlapping segments

n_seeds = 1000
n_modes=10
n_samples = 50 #bootstrapping samples
range_delays = np.arange(1,21)
size = 10000

labels = ma.array(cl.kmeans_knn_partition(traj_matrix,n_seeds),dtype=int)

delay_range = np.arange(1,20,1)
n_modes=5
ts_traj = np.zeros((len(delay_range),n_modes))
for kd,delay in enumerate(delay_range):
    P = op_calc.transition_matrix(labels,delay)
    R = op_calc.get_reversible_transition_matrix(P)
    ts_traj[kd,:] = op_calc.compute_tscales(R,delay,dt,k=n_modes+1)
    print(delay)
# -

plt.plot(delay_range*dt,ts_traj)
plt.show()

delay = 12
print(delay*dt)
P = op_calc.transition_matrix(labels,delay)
R = op_calc.get_reversible_transition_matrix(P)
eigvals,eigvecs = op_calc.sorted_spectrum(R,k=3)
print(eigvals)

X=eigvecs[:,1].real

plt.figure(figsize=(10,7))
color_abs = np.max(np.abs(X))
plt.scatter(traj_matrix[:,0],traj_matrix[:,-1],c=X[labels],cmap='seismic',s=1,vmin=-color_abs,vmax=color_abs)
# plt.xticks(range(-20,21,5))
# plt.ticks()
# plt.savefig('Phi_2_Lorenz.png')
plt.show()



StarcoderdataJupyterScriptsDedupFiltered