TIR.py

#0. software packages

#os
import os
import os.path
import sys
from time import time
from datetime import datetime

#basics
import numpy as np
from numpy.random import random
import pandas as pd
import math
import scipy
import pickle
from copy import deepcopy

#plotting
import matplotlib.pyplot as plt
from matplotlib.ticker import StrMethodFormatter
from matplotlib import rcParams
import matplotlib as mpl
from matplotlib  import cm
from matplotlib import colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.offsetbox import (TextArea, DrawingArea, OffsetImage,AnnotationBbox)


#training
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn import manifold
from scipy.linalg import orthogonal_procrustes,eigh
from scipy.stats import linregress
from sklearn.metrics.pairwise import euclidean_distances
from sklearn.metrics import mean_squared_error

# torch is imported lazily inside lift_returns_ann and the legacy ANN
# trainers so that TIR.py is usable for the non-ANN pipeline (delay vectors,
# diffusion maps, bootstrap CIs, Kabsch, panel-B helpers) on environments
# without torch installed.
try:
    from torch import nn
    from torch.utils.data import Dataset, TensorDataset, DataLoader
    import torch.nn.parallel
    from torch.optim import Adam
    from torch.utils.data.dataset import random_split
    _HAS_TORCH = True
except ImportError:
    _HAS_TORCH = False

#Entropies
####################################################################################################################
# 1.1: Shanon Entropies
def shan_ts(ser,b):
    """
    Compute the Shannon entropy of a time series.

    Args:
        ser: The time series (array).
        b: The number of bins to use for the histogram (int).

    Returns:
        float: The Shannon entropy of the time series.
    """
    c = np.histogram(ser,b)[0]
    c_normalized = c / float(np.sum(c))
    c_normalized = c_normalized[np.nonzero(c_normalized)]
    H = -sum(c_normalized* np.log2(c_normalized))  
    return H

def shan(c):
    """
    Compute the Shannon entropy from histogram counts.

    Args:
        c: The time series histogram (array).

    Returns:
        float: The Shannon entropy of the time series.
    """
    c_normalized = c / float(np.sum(c))
    c_normalized = c_normalized[np.nonzero(c_normalized)]
    H = -sum(c_normalized* np.log2(c_normalized))  
    return H
####################################################################################################################
# 1.2: Autocorrelation and Mutual Information
def AC(ser,ts,dim=5,maxval=1000,tol=10,cut=0.2,save_loc='N',save=False,plots=False, method='cutoff'):
    """
    Compute the autocorrelation of a time series.

    Args:
        ser: The time series (array).
        ts: The time series name (string)
        dim: The dimensionality of the time series (int).
        maxval: The maximum value of tau to test (int).
        tol: The time delay step size (int).
        cut: The cutoff value determining at which time delay to stop, therby selecting the appropriate tau (float).
        save_loc: The location to save the image (string).
        save: A boolean to save the time series (boolean).
        plots: A boolean to plot the plots (boolean).
        method: The method to determine tau ('cutoff' or 'minimum') (string).

    Returns:
        int: The optimal time delay of the time series.
    """
    
    # make an array of tested time delays from 0 to maxval with step size tol
    tau_arr = np.arange(0,maxval,tol)
    # define the mean, std and deviation from mean for each value in TS
    mu, sigma= np.mean(ser), np.std(ser)
    mz =  ser-mu
    #Store autocorrelations in array
    AC_arr = []
    #compute AC at each value of tau, t
    for t in tau_arr:
        val = np.mean(np.multiply(mz[0:len(mz)-t],mz[t:len(mz)]))/(sigma**2) 
        AC_arr.append(val)
    #Determine cutoff by computing 
    if method == 'cutoff':
        for tau_i in range(len(AC_arr)):
            if AC_arr[tau_i] <= cut:
                AC_length=tau_arr[tau_i]
                break
            #choose last val if no other option    
            elif tau_i == len(AC_arr)-1:
                AC_length = int(tau_arr[tau_i]*tol)

    
    elif method == 'minimum':
        # Determine tau by finding the minimum autocorrelation value

        #step 1: polyfit to get rid of noise
        x = np.arange(len(AC_arr))
        poly_coeffs = np.polyfit(x, AC_arr, 6)
        poly_fit = np.polyval(poly_coeffs, x)

        #step 2: find minimum
        for jkl in range(1, len(poly_fit) - 1):
            if poly_fit[jkl-1]>poly_fit[jkl] and poly_fit[jkl+1]>poly_fit[jkl]:   
                AC_length=int(jkl*tol)
                break
            else:
                AC_length=int(np.argmin(AC_arr)*tol)
    print(AC_length)
    #plot AC 
    if plots==True:
        fig, ax = plt.subplots()

        ax.plot(tau_arr, AC_arr, 'b-', marker='o')
        ax.plot([tau_arr[0],tau_arr[-1]], [0,0], 'k:')
        ax.set_xlabel("delay time, tau (steps)")
        ax.set_ylabel("autocorrelation, (arb. units)")
        ax.figure.savefig("%s/AC_%s.pdf" %(save_loc,ts))
        plt.draw()
     #save param
    print('Suggested time delay', AC_length)
    if save==True:
        np.savez('%s/AC_%s_tau%d_dim%d.npz' %(save_loc,ts,AC_length,dim),AC=AC_length)
    plt.show()
    plt.close()
    return AC_length
####################################################################################################################
# 1.3: Mutual Information
def MI(x,y,b=100):
    """
    Compute the autocorrelation of a time series.

    Args:
        x: The first time series (array).
        y: The second time series (array).
        b: The number of bins to use for the histogram (int).

    Returns:
        float: The mutual information between two time series
    """
    c_x=np.histogram(x,bins=b)[0]
    c_y=np.histogram(y,bins=b)[0]
    c_xy=np.histogram2d(x,y,bins=b)[0]
    return shan(c_x)+shan(c_y) - shan(c_xy)
####################################################################################################################
# 1.4: Draw Mutual Information Curve
def MI_curve(ser1,ser2,ts,b=100,maxval=1000,tol=1,save=False,save_loc='N',cut=0.2):
    """
    Compute the MI curves for two time series

    Args:
        ser1: The first time series (array).
        ser2: The second time series  (array).
        ts: The time series name (string).
        b: The number of bins to use for the histogram (int).
        maxval: The maximum value of tau to test (int).
        tol: The time delay step size (int).
        save: A boolean to save the time series (boolean).
        save_loc: The location to save the image (string).
        cut: The cutoff value determining at which time delay to stop, therby selecting the appropriate tau (float).

    Returns:
        float: The mutual information between two time series 
        int: The optimal time delay of the time series.
    """

    MI_v = []
    for iii in range(0,maxval,tol):
        x = ser1[0:len(ser1)-iii]
        y = ser2[iii:len(ser1)]
        c_x=np.histogram(x,bins=b)[0]
        c_y=np.histogram(y,bins=b)[0]
        c_xy=np.histogram2d(x,y,bins=b)[0]
        MI_v.append(shan(c_x)+shan(c_y) - shan(c_xy))

    
    #step 1: polyfit to get rid of noise
    x = np.arange(len(MI_v))
    poly_coeffs = np.polyfit(x, MI_v, 6)
    poly_fit = np.polyval(poly_coeffs, x)

    #step 2: find minimum
    for jkl in range(1, len(poly_fit) - 1):
        if poly_fit[jkl-1]>poly_fit[jkl] and poly_fit[jkl+1]>poly_fit[jkl]:  
            MI_length=int(jkl*tol)
            break
        else:
            MI_length=int(np.argmin(MI_v)*tol)

    plt.scatter(range(len(MI_v)),MI_v,linewidth=3)
    plt.xlabel('Time Delay (time steps)',fontsize=15)
    plt.ylabel('Mutual Information (arb. units)',fontsize=15)
    plt.show()

    if save==True:
        plt.savefig('%s/MI_%s.pdf' %(save_loc,ts))
    plt.close()

    return MI_v,MI_length

####################################################################################################################
# 1.5: Optimal embedding dimensionality (E1d) 
def E1d_cal(ser,ts,stop,stride=10,tau=1,dim=5,dMax=25,start=0,tol=0.98,save_loc='N',save=False,plots=False):
    """
    Estimate the optimal embedding dimensionality of a time series using the E1d method of Cao (1997).

    Args:
        ser: The time series (array).
        ts: The time series name (string)
        stop: The stopping point for the E1d analysis (int).
        stride: The stride for the E1d analysis (int).
        tau: The time delay for the E1d analysis (int).
        dim: The maximum dimensionality for the E1d analysis (int).
        dMax: The maximum dimensionality for the E1d analysis (int).
        start: The starting point for the E1d analysis (int).
        tol: The tolerance for the E1d analysis (float).
        save_loc: The location to save the image (string).
        save: A boolean to save the time series (boolean).
        plots: A boolean to plot the plots (boolean).

    Returns:
        int: The optimal embedding dimensionality of the time series.
    """
    # plotting delay vectors
    if plots==True:
       
        #construct delay vectors via Takens Theorem
        delayVecs = np.reshape(ser[0:len(ser)-(dim-1)*tau],(len(ser)-(dim-1)*tau,1))
        for ii in range(1,dim):
            delayVecs = np.concatenate((delayVecs,np.reshape(ser[ii*tau:len(ser)-(dim-1)*tau+ii*tau],(len(ser)-(dim-1)*tau,1))), axis=1)
        print('nVecs = %d' % delayVecs.shape[0])
        plot_stride = 10
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        i,j,k=0,2,4
        ax.scatter(delayVecs[::plot_stride,i], delayVecs[::plot_stride,j], delayVecs[::plot_stride,k])
        ax.set_xlabel('delayVec_' + str(i))
        ax.set_ylabel('delayVec_' + str(j))
        ax.set_zlabel('delayVec_' + str(k))
        plt.draw()
        plt.show()
        plt.close()
    
    # Performing E1d analysis using specified value of tau
    #initialize Ed and define max dim
    Ed=[]
    #build Takens delay vectors for each dimension
    for d in range(1,dMax):
        # building d-dim delay embedding of len(h2t)-(d-1)*tau delay points with delay time of tau (in steps)
        delayVecs = np.reshape(ser[0:len(ser)-(d-1)*tau],(len(ser)-(d-1)*tau,1))
        for ii in range(1,d):
            delayVecs = np.concatenate((delayVecs,np.reshape(ser[ii*tau:len(ser)-(d-1)*tau+ii*tau],(len(ser)-(d-1)*tau,1))), axis=1)
       
        # subsampling -  check if start and stop are within bounds
        if start<0 | stop>delayVecs.shape[0]:
            print(stop,delayVecs.shape[0])
            print('ERROR: subsampling illegal points')
            break
        delayVecs_ss = delayVecs[start:stop:stride,:]
    
        # construct pairwise distance matrix between delay vectors pairwise distances between delay vectors
        P = euclidean_distances(delayVecs_ss,delayVecs_ss)
        idx = np.arange(0,delayVecs_ss.shape[0],1)

        # computing a_id and a_i(d+1) for each point
        if d >= 2:
            Ed.append(np.mean(np.divide(P[idx,idxNN_d],deltaNN_d)))

        # identifying non-zero nearest neighbor of each point
        P[np.where(P==0)]=np.inf
        idxNN_d = np.argmin(P, axis=1)
        deltaNN_d = P[idx,idxNN_d]

    # computing E1d
    Ed = np.array(Ed)
    E1d = np.divide( Ed[1:len(Ed)], Ed[0:len(Ed)-1] )
    #select E1d that is odd and explains more than tol of variance that is also odd (oddness only matters with temporal symmetry)
    for k in range(len(E1d)):
        if k//2!=0 and E1d[k] >=tol:
            E1dmax=k
            break
        #set to last value if no other option; report
        else:
            E1dmax=k
            print('Increase dMax or reduce tolerance, no suitable E1d found')
    
    if plots==True:
        #make a plot
        fig, ax = plt.subplots()
        ax.plot(np.arange(2,len(E1d)+1), E1d[1::], 'k-', marker='o')
        ax.set_xlabel("delay dim")
        ax.set_ylabel("E1d")
        ax.figure.savefig("%s/E1d_%s_tau%d_dim%d.pdf" %(save_loc,ts,tau,dim))
        plt.draw()
    if save==True:
        np.savez('%s/E1d_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim),E1d=E1dmax)
    print('Suggested dimensionality: ', E1dmax )
    return E1dmax

####################################################################################################################
#Scaling Functions
#2.1 Standard Scaling for Time Series
def scalearr(ser,ts='N',dim=5,tau=1,save_loc='N',save=False):
    """
    Scale a time series to the range [0,1].

    Args:
        ser: The time series (array).
        ts: The time series name (string)
        dim: The dimensionality of the time series (int).
        tau: The time delay of the time series (int).
        save_loc: The location to save the scaling (string).
        save: A boolean to save the scaling (boolean).

    Returns:
        array: The scaled time series (array).
    """
    
    ser = ser.reshape(-1, 1) 
    scaler = MinMaxScaler(feature_range=(0,1)) 
    scaled_ts = scaler.fit_transform(ser)
    # saving scaling associated with model
    if save==True:
        D = {}
        D['zscale'] = scaler
        file = open('%s/scaled_%s.pkl' %(save_loc,ts), 'wb')
        pickle.dump(D, file)
        file.close()
        np.savez('%s/scaled_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim),scaled_ts=scaled_ts)
    return scaled_ts

####################################################################################################################
# 2.2: Standard Scaling for Manifold
def unscalearr(Z_pred,load_loc,ts='N',dim=5,tau=1,save_loc='N',save=False):
    """
    Unscale z-scaled manifold

    Args:
        Z_pred: The z-scaled manifold (array).
        load_loc: The location to load the scaling (string).
        ts: The time series name (string).
        dim: The dimensionality of the time series (int).
        tau: The time delay of the time series (int).
        save_loc: The location to save the scaling (string).
        save: A boolean to save the scaling (boolean).

    Returns:
        array: The unscaled manifold.
    """
    D = pickle.load(open(load_loc, 'rb'))
    zscale_aa = D['zscale']
    z_pred = zscale_aa.inverse_transform(Z_pred)
    if save==True:
        np.savez('%s/unscaled_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim),z_pred=z_pred)
    return z_pred

####################################################################################################################
# 2.3: Standard Scaling for Pandas Dataframe
def scalepd(TS,ts,dim=5,tau=1,save_loc='N',save=False):
    """
    Scale a time series to the range [0,1] from a pandas dataframe.

    Args:
        TS: The time series (dataframe).
        ts: The time series name (string).
        dim: The dimensionality of the time series as computed by E1d (int).
        tau: The time delay of the time series (int).
        save_loc: The location to save the scaling (string).
        save: A boolean to save the scaling (boolean).

    Returns:
        array: The scaled time series.
    """

    ser = np.array(TS.loc[:,ts]).reshape(-1, 1)
    scaler = MinMaxScaler(feature_range=(0,1)) 
    scaled_ts = scaler.fit_transform(ser)
    # saving scaling associated with model
    if save==True:
        D = {}
        D['zscale'] = scaler
        file = open('%s/scaled_%s.pkl' %(save_loc,ts), 'wb')
        pickle.dump(D, file)
        file.close()
        np.savez('%s/scaled_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim),scaled_ts=scaled_ts)
    return scaled_ts


# Delay Vector construction
####################################################################################################################
# 3. Structure data (singlevairate, multivariate, multitemporal)

def make_dvecs(ser,ts,tau,dim,save=False,save_loc='N',delaytau=0,mode='single',scaling='None'):
    """
    Make Takens delay vectors.

    Args:
        ser: The time series (array) or (dataFrame) if mode is 'multivariate'.
        ts: The time series name (string)
        tau: The time delay of the time series (int). If mode is 'multivariate' it is an array of delays.
        dim: The dimensionality of the time series (int). If mode is 'multivariate' it is an array of delays.
        save_loc: The location to save the delay vectors (string).
        save: A boolean to save the delay vectors (boolean).
        delaytau: Eliminate the initial x values of the time series for equilibration (int).
        mode: The mode of the time series ('single', 'multivariate', 'multitemporal') (string).
        scaling: The scaling of the time series ('None', 'standard', 'return') (string).

    Returns:
        array: Single variate Takens delay vectors of shape (len(ser)-(dim-1)*tau, dim).
    """
    if mode=='single': 
        ser = ser[delaytau::]
        print(len(ser))
        delayVecs = np.reshape(ser[0:len(ser)-(dim-1)*tau],(len(ser)-(dim-1)*tau,1))
        for ii in range(1,dim):
            delayVecs = np.concatenate((delayVecs,np.reshape(ser[ii*tau:len(ser)-(dim-1)*tau+ii*tau],(len(ser)-(dim-1)*tau,1))), axis=1)
        

        #standard deviation scaling of delay vectors
        if scaling=='standard':
            for kk in range(delayVecs.shape[0]):
                delayVecs[kk,:] = (delayVecs[kk,:] - np.mean(delayVecs[kk,:]))/np.std(delayVecs[kk,:])

        #returns scaling of delay vectors
        elif scaling=='return':
            for kk in range(delayVecs.shape[0]):
                # each row has each position subtracted by the first position for return and divided by the mean value of the row
                mval = np.mean(delayVecs[kk,:])
                delayVecs[kk,:] = np.divide((np.subtract(delayVecs[kk,:],delayVecs[kk,0])),mval)


    if mode=='multivariate': #let ser is a dataframe
        iii = 0
        #reshape the portion of ser accessible to takens of dim*tau
        for serz in ser.columns:
            if isinstance(tau,int)==False:
                tauv = tau[iii]
                dimv = dim[iii]
        
            #define which ts to use
            serp = np.array(ser[serz])[delaytau:]
            #append all vals to each element in dvecs as in Cao 1998
            if iii==0:
                delayVecs = [[] for x in range(len(serp)-(dimv-1)*tauv)]

            dvecs_mult = np.reshape(serp[0:len(serp)-(dimv-1)*tauv],(len(serp)-(dimv-1)*tauv,1)) 
            for ii in range(1,dimv):
                dvecs_mult = np.concatenate((dvecs_mult,np.reshape(serp[ii*tauv:len(serp)-(dimv-1)*tauv+ii*tauv],(len(serp)-(dimv-1)*tauv,1))), axis=1)   

            #standard deviation scaling of delay vectors
            if scaling=='standard':
                for kk in range(delayVecs.shape[0]):
                    dvecs_mult[kk,:] = (dvecs_mult[kk,:] - np.mean(dvecs_mult[kk,:]))/np.std(dvecs_mult[kk,:])

            #returns scaling of delay vectors
            elif scaling=='return':
                for kk in range(dvecs_mult.shape[0]):
                    # each row has each position subtracted by the first position for return and divided by the mean value of the row
                    mval = np.mean(dvecs_mult[kk,:])
                    dvecs_mult[kk,:] = np.divide((np.subtract(dvecs_mult[kk,:],dvecs_mult[kk,0])),mval)
            #combine delay vectors
            if iii>=1:
                mlen= int(min(len(dvecs_mult),len(delayVecs)))
            else:
                mlen = len(dvecs_mult)
            delayVecs= np.hstack((delayVecs[0:mlen],dvecs_mult[0:mlen]))
            iii+=1



    if mode=='multitemporal': #let ser is an array, tau is an array of delays, dim is dimension
         #max dvec size:
        ser = ser[delaytau:]
        print('max size is nVecs = %d' %(len(ser)-(dim-1)*max(tau)))
        #initialize
        delayVecs = [[] for x in range(len(ser)-(dim-1)*tau[0])]
        iii = 0
        #reshape the portion of ser accessible to takens of dim*tau
        for tauv in tau:
        #   dvecs_mult = []
            #append all vals to each element in dvecs as in Cao 1998
            #make dvecs for each given tau
            dvecs_mult = np.reshape(ser[0:len(ser)-(dim-1)*tauv],(len(ser)-(dim-1)*tauv,1)) 
            #make dvecs for each observation
            for ii in range(1,dim):
                dvecs_mult = np.concatenate((dvecs_mult,np.reshape(ser[ii*tauv:len(ser)-(dim-1)*tauv+ii*tauv],(len(ser)-(dim-1)*tauv,1))), axis=1)  
            
            #standard deviation scaling of delay vectors
            if scaling=='standard':
                for kk in range(delayVecs.shape[0]):
                    dvecs_mult[kk,:] = (dvecs_mult[kk,:] - np.mean(dvecs_mult[kk,:]))/np.std(dvecs_mult[kk,:])

            #returns scaling of delay vectors
            elif scaling=='return':
                for kk in range(dvecs_mult.shape[0]):
                    # each row has each position subtracted by the first position for return and divided by the mean value of the row
                    mval = np.mean(dvecs_mult[kk,:])
                    dvecs_mult[kk,:] = np.divide((np.subtract(dvecs_mult[kk,:],dvecs_mult[kk,0])),mval)
            #compute min length of dvec to make sure we can concatenate
            if iii>=1:
                mlen= int(min(len(dvecs_mult),len(delayVecs)))
            #set to whatever first dvec is otherwise
            else:
                mlen = len(dvecs_mult)
            #hstack concatenated dvecs as per Cao 1998
            delayVecs= np.hstack((delayVecs[0:mlen],dvecs_mult[0:mlen]))
            iii+=1

    # saving delay vectors
    if save==True:
        np.savez('%s/dvecs_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim), dvecs=delayVecs)
    return delayVecs
####################################################################################################################
#Diffusion map
#4.1: Pairwise distance matrix
def make_P(ts,dvecs,train,tau,dim,save=False,save_loc='N',issymmetric=False,distmetric='euclidean'):
    """
    Compute the distance matrix of the time series.

    Args:
        ts: The time series name (string).
        dvecs: The delay vectors of the time series (array).
        train: The training data (array).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        save_loc: The location to save the diffusion map (string).
        save: A boolean to save the diffusion map (boolean).
        issymmetric: A boolean to determine if the diffusion map is symmetric (boolean).
        distmetric: The distance metric to use ('euclidean', 'manhattan', 'cosine') (string).

    Returns:    
        array: The pairwise distance matrix of the time series.
    """ 
    
    dvecs_ss = deepcopy(dvecs[train])
    # euclidean distance between two vectors
    if distmetric=='euclidean':
        if issymmetric==False:
            P = euclidean_distances(dvecs_ss,dvecs_ss)
        else:
            dvecs_ss_flip = np.flip(dvecs_ss,axis=0)
            P = np.minimum( euclidean_distances(dvecs_ss,dvecs_ss), euclidean_distances(dvecs_ss_flip,dvecs_ss) )

    # manhattan distance between two vectors
    elif distmetric=='manhattan':
        if issymmetric==False:
            P = manhattan_distances(dvecs_ss,dvecs_ss)
        else:
            dvecs_ss_flip = np.flip(dvecs_ss,axis=0)
            P = np.minimum( manhattan_distances(dvecs_ss,dvecs_ss), manhattan_distances(dvecs_ss_flip,dvecs_ss) )
    #cosine similarity of two vectors

    elif distmetric=='cosine':
        if issymmetric==False:
            P = cosine_distances(dvecs_ss,dvecs_ss)
        else:
            dvecs_ss_flip = np.flip(dvecs_ss,axis=0)
            P = np.minimum( cosine_distances(dvecs_ss,dvecs_ss), cosine_distances(dvecs_ss_flip,dvecs_ss) )
    if save==True:
        np.savez('%s/P_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim), P=P)
    return P
####################################################################################################################
# 4.2: Gaussian Kernel Width Search
def eps_scan(ts,tau,dim,P,dvecs,train,ll=-10,ul=10,res=0.5,tol=0.85,save=False,save_loc='N',plots=False,return_all=False,nEvals=25,i_dim=7,method='linear',fit=2,linlim='Minlinlim',clim=20,mineps=-1,lamulim=0.9,lamlim=0.3,lamllim=0.1):
    """
    Compute the optimal epsilon value for the diffusion map.

    Choice for EPS can be made such that subsequent dMap has dimension equal to E1d analysis result OR
    condition listed on page 6 of DOI: 10.1109/TPAMI.2008.76 (algorithm 3.1)for kernel selection as default
    note also that we choose the inflection point as the center of the linear region of the log(sumA) vs log(eps) curve wheras the 
    condition listed in the paper is that epsilon be chosen from the linear region [epsilon_min,epsilon_max].


    Args:
        ts: The time series name (string).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        P: The distance matrix of the time series (array).
        dvecs: The delay vectors of the time series (array).
        train: The training data (array).
        ll: The lower limit of the epsilon scan (float).
        ul: The upper limit of the epsilon scan (float).
        res: The resolution of the epsilon scan (float).
        tol: The tolerance for the epsilon scan (float).
        save_loc: The location to save the epsilon scan (string).
        save: A boolean to save the epsilon scan  (boolean).
        plots: A boolean to plot the epsilon scan (boolean).
        return_all: A boolean to return all epsilon values (boolean).
        nEvals: The number of eigenvalues to compute (int).
        i_dim: The intrinsic dimensionality of the data computed from E1d (int).
        method: The method to determine epsilon ('manual', 'linear', 'intrinsic') (string).
        fit: Numbers of points to use for the fit for the epsilon scan (int).
        linlim: The limit for the epsilon scan ('Maxlinlim', 'Minlinlim') (string).
        clim: The number of attempts to find the epsilon value (int).
        mineps: The minimum epsilon value for the epsilon scan (float).
        lamulim: The upper limit for the epsilon scan (float).
        lamlim: The limit for the epsilon scan (float).
        lamllim: The lower limit for the epsilon scan (float).
    
    Returns:
        float: The optimal epsilon value for the diffusion map.
    """ 

    N = P.shape[0]
    log_eps_array = np.arange(ll,ul,res)
    sumA = np.empty_like(log_eps_array)
    for i in range(len(log_eps_array)):
        eps = np.exp(log_eps_array[i])
        sumA[i] = np.sum(np.exp(-P**2/(2*eps)),axis=None)

    if method=='manual':  
        for ij in range(len(sumA)):
            i = np.log(sumA[ij])
            if i >=  (np.log(sumA[0])+np.log(sumA[-1]))/2: #(np.log(N) - np.log(N)*tol): 
                epsval = log_eps_array[ij]
                break
            elif ij==len(sumA)-1:
                print('no epsval found, investigate range choice')
                print('See DOI: 10.1109/TPAMI.2008.76 P7-9 for kernel selection criteria')
                epsval = ul

                #compare with L method
        epsval = log_eps_array[indx] 
        P = make_P(ts,dvecs,train,tau,dim)
        #compute dMap to get eval spectrum for restricted dataset
        lamb,psi,eps = dMap(ts,tau,dim,P,epsval,nEvals)
        tdim = L_method(lamb,plots=True)
        print('Implied dimensionality: ', tdim)

    #Verlet 2nd derivative
    elif method=='linear':
        SD = []
        for ij in range(1,len(sumA)-1):
            # take first derivative
            #estimate of SD
            SD.append((np.log(sumA[1+ij])+np.log(sumA[ij-1])-2*np.log(sumA[ij]))/res**2)
        plt.scatter(range(len(SD)),SD)
        plt.xlabel('eps index')
        plt.ylabel('second derivative')
        plt.xticks(np.arange(0,len(SD),1),log_eps_array[1:-1])
        plt.title('Second Derivative of log(sumA) vs log(eps)')
        plt.show()
        plt.close()
            #check sign change
        #iterate over len SD to find inflection point
        #why argmax? In the case where there are two linear regions we want the region of largest difference
        print(np.argmax(SD),SD[np.argmax(SD)] )
        for ij in range(np.argmax(SD)-2,len(SD)):
            # double check if second derivative is positive and then negative for confirmed inflection point
            if SD[ij-1]>=0 and SD[ij]<=0:
                indx = ij+1-fit #why? +1 for the SD  because we want to start on the far left of the linear region and work our way to correct dim
                epsval = log_eps_array[indx]
                break
            else:
                indx = ij

        if linlim=='Maxlinlim':
            epsval = log_eps_array[indx+fit]
        elif linlim=='Minlinlim':
            epsval = log_eps_array[indx-fit]
        #compare with L method
        epsval = log_eps_array[indx] 
        P = make_P(ts,dvecs,train,tau,dim)
        #compute dMap to get eval spectrum for restricted dataset
        lamb,psi,eps = dMap(ts,tau,dim,P,epsval,nEvals)
        tdim = L_method(lamb,plots=True)
        print('Implied dimensionality: ', tdim)
        

    #match the intrinsic dimensionality of the data as per DOI: 10.1109/TPAMI.2008.76 manifold learning choice
    # L method of Salvador and Chan - https://cs.fit.edu/~pkc/papers/ictai04salvador.pdf 
    elif method=='intrinsic':
        #second derivative: find point of inflection where rate of change of slope goes from positive to negative
        SD = []
        for ij in range(1,len(sumA)-1):
            # take first derivative
            #estimate of SD
            SD.append((np.log(sumA[1+ij])+np.log(sumA[ij-1])-2*np.log(sumA[ij]))/res**2)
        plt.figure(figsize=(20,10))
        plt.scatter(range(len(SD)),SD)
        plt.xlabel('eps index')
        plt.ylabel('second derivative')
        plt.xticks(np.arange(0,len(SD),1),log_eps_array[1:-1])  
        plt.title('Second Derivative of log(sumA) vs log(eps)')
        plt.show()
        plt.close()
        #iterate over len SD to find inflection point
        #why argmax? In the case where there are two linear regions we want the region of largest difference
        for ij in range(1,len(SD)): 
            # double check if second derivative is positive and then negative for confirmed inflection point
            if SD[ij-1]>=0 and SD[ij]<=0:
                indx = ij+1-fit #why? +1 for the SD  because we want to start on the far left of the linear region and work our way to correct dim
    
                break

            
            #if we never get to inflection point:
            elif ij==len(SD)-1:
                print('No epsval found, investigate range choice. Increase range and/or resolution')
                print('See DOI: 10.1109/TPAMI.2008.76 P7-9 for kernel selection criteria')
                print('Displaying spectrum')
                fig, ax = plt.subplots()
                ax.scatter(log_eps_array, np.log(sumA))
                ax.plot([log_eps_array[0],log_eps_array[-1]],[np.log(N**2),np.log(N**2)],'k:')
                ax.plot([log_eps_array[0],log_eps_array[-1]],[np.log(N),np.log(N)],'k:')
                ax.set_xlabel("log(eps)")
                ax.set_ylabel("log(sumA)")
                plt.show()
                plt.close()
                indx = ij
            


        #first guess at correct epsilon; we use - fit to ensure we begin in the over saturated region
        epsval = log_eps_array[indx] 
        if epsval <mineps:
            epsval = mineps
        P = make_P(ts,dvecs,train,tau,dim)
        #compute dMap to get eval spectrum for restricted dataset
        lamb,psi,eps = dMap(ts,tau,dim,P,epsval,nEvals)
        tdim = L_method(lamb,plots=True)
        #we check if the extracted dimensionality is smaller than the intrinsic dim from E1d
        #if yes, we augment the gaussian kernel estimation and recompute the eigenvalue spectrum
        #it is recommended to use plots=True to see the spectrum and the linear fit
        
        val = 1 
        mem = [0]
        #while tdim - i_dim > 1 and np.mean(lamb[1:3])>0.9 and lamb[1]>0.3 and val < 20 or np.mean(lamb[10:15])>0.3: #should not be converged. if converged, stop mem[-1]!=mem[-2]
        while val < clim:
            if tdim - i_dim > 1 and np.mean(lamb[1:3])>lamulim and lamb[1]>lamlim or np.mean(lamb[12:20])>lamllim:
        
                try: # to deal with convergence issues leading to linalg error, we try and
                    epsval = log_eps_array[indx+val] 
                    if log_eps_array[indx+val]<=mineps:
                        epsval = mineps
                        #accelerate sampling
                    if np.mean(lamb[1:15])>lamulim:
                        val +=2
                        epsval = log_eps_array[indx+val]
                    print('epsval we are trying: ', epsval)
                    P = make_P(ts,dvecs,train,tau,dim)
                    lamb,psi,eps = dMap(ts,tau,dim,P,epsval,nEvals)
                    tdim = L_method(lamb,plots)
                    mem.append(tdim)
                    val+=1

                    print('max eval is: ', lamb[1])
                except:
                    print('skipped epsval: ', log_eps_array[indx+val] )
                    val +=1
                    continue
                    #pass
            else:
                print('Convergence achieved, implied dimensionality: ', tdim)
                print('mean eval is: ', np.mean(lamb[1:15]))
                print(mem)
                break
    
         
    if plots==True:
        fig, ax = plt.subplots()
        ax.scatter(log_eps_array, np.log(sumA))
        ax.plot([log_eps_array[0],log_eps_array[-1]],[np.log(N**2),np.log(N**2)],'k:')
        ax.plot([log_eps_array[0],log_eps_array[-1]],[np.log(N),np.log(N)],'k:')
        ax.set_xlabel("log(eps)")
        ax.set_ylabel("log(sumA)")
        fig.savefig('%s/eps_%s_tau%d_dim%d.png' %(save_loc,ts,tau,dim), dpi=300)

       #straight line fit: - let slope be the first derivative at the inflection point
        slope = (np.log(sumA[indx+fit])-np.log(sumA[indx-fit]))/(2*fit*res)
        intersec = np.log(sumA[indx])-slope*log_eps_array[indx]
        ax.plot(log_eps_array[indx-fit:indx+fit+1], slope*log_eps_array[indx-fit:indx+fit+1]+intersec,color='red')       

    if save==True:
        np.savez('%s/eps_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim), epsval=epsval)

    # DOI: 10.1109/TPAMI.2008.76
    #apply the linearity condition
    return epsval

####################################################################################################################
# 4.3: Diffusion Map
def dMap(ts,tau,dim,P,epsval,nEvals,save=False,save_loc='/N',plots=False,plot_stride=100):
    """
    Compute the diffusion map from a pairwise distance matrix.

    Args:
        ts: The time series name (string).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        P: The distance matrix of the time series (array).
        epsval: The optimal epsilon value for the diffusion map (float).
        nEvals: The number of eigenvalues to compute (int).
        save_loc: The location to save the diffusion map (string).
        save: A boolean to save the diffusion map (boolean).
        plots: A boolean to plot the diffusion map (boolean).
        plot_stride: The stride of the diffusion map plot (int).  

    Returns:
        array: The eigenvalues of the diffusion map (array).
        array: The eigenvectors of the diffusion map (array).
        float: The optimal epsilon value for the diffusion map (float).
    """

    eps = np.exp(epsval)
    P = np.exp(-P**2/(2*eps))
    D = np.sum(P,axis=1)
    P = np.matmul(np.matmul(np.diag(D**(-0.5)),P),np.diag(D**(-0.5))) # constructing Ms as symmetric matrix for diagonalization

    lamb, psi = eigh(P,subset_by_index=(P.shape[0]-nEvals,P.shape[0]-1)) # scipy eigh to specify only computation of leading evals
    print(lamb.shape)
    psi = np.matmul(np.diag(D**(-0.5)),psi) # converting evecs of Ms to evecs of M; M and Ms share evals
    idx_sort = np.flip(np.argsort(lamb))
    lamb,psi = lamb[idx_sort],psi[:,idx_sort]

    if save==True:
        np.savez('%s/map_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim), lamb=lamb, psi=psi)

    if plots==True:
        #plot
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.bar(np.arange(1,len(lamb)), lamb[1:])
        ax.set_xlabel('Eigenvalue index',fontsize=15)
        ax.set_ylabel('Eigenvalues',fontsize=15)
        ax.set_title('Eigenvalue Spectrum',fontsize=15)
        plt.xticks(np.arange(1, len(lamb), step=1))
        plt.tight_layout()
        fig.savefig('%s/evals_eps_%s_tau%d_dim%d.png' %(save_loc,ts,tau,dim), dpi=300)
        plt.show()
        plt.close()

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        i,j,k=1,2,3
        ax.scatter(psi[::plot_stride,i], psi[::plot_stride,j], psi[::plot_stride,k])
        ax.set_xlabel(r'$\Psi_%s$' %str(i),fontsize=15)
        ax.set_ylabel(r'$\Psi_%s$' %str(j),fontsize=15)
        ax.set_zlabel(r'$\Psi_%s$' %str(k),fontsize=15)
        ax.set_title('Diffusion Map Embedding',fontsize=15)
        plt.tight_layout()
        plt.draw()
        plt.show()
        plt.close()
        plt.savefig('%s/dMap_%s_tau%d_dim%d.pdf' %(save_loc,ts,tau,dim))
   
    return lamb,psi,eps

#####################################################################################################################
# 4.4: Intrinsic Dimensionality
def L_method(lamb,plots=False):
    """
    Estimate the intrinsic dimensionality of the data from the eigenvalue spectrum

    Args:
        lamb: The eigenvalues of the diffusion map (array).
        plots: A boolean to plot the intrinsic dimensionality (boolean).

    Returns:
        int: The intrinsic dimensionality of the data (int).
    """ 
    
    TRMSE = []
    lamb = lamb[1:]
    lamb_indx = np.linspace(1,len(lamb),len(lamb))#[0:-1]
    for ij in range(1,len(lamb)-1):
        l_fit1 = np.polyfit(lamb_indx[0:ij+1],lamb[0:ij+1],1)
        l_fit2 = np.polyfit(lamb_indx[ij:],lamb[ij:],1)
        #compute RMSEs for each fit with the data
        RMSE1 =(lamb[0:ij+1]-np.polyval(l_fit1,lamb_indx[0:ij+1]))
        RMSE2 = (lamb[ij:]-np.polyval(l_fit2,lamb_indx[ij:]))
        RMSE3 = np.array(list(RMSE1)+list(RMSE2))
        TRMSE.append(np.sqrt(np.mean(RMSE3**2)))

    #find minimum TRMSE
    #why +1? - because we are indexing from 1. Still need to add 1 for intrinsic dim in labeling
    true_dim = int(lamb_indx[np.argmin(TRMSE)])
    if plots==True:
        l_fit1 = np.polyfit(lamb_indx[0:true_dim+1],lamb[0:true_dim+1],1)
        l_fit2 = np.polyfit(lamb_indx[true_dim:],lamb[true_dim:],1)
        l1 = np.polyval(l_fit1,lamb_indx)
        l2 = np.polyval(l_fit2,lamb_indx)
        l1 = l1[l1>-0.2]
        #plot
        fig = plt.figure(figsize=(10,5))
        ax = fig.add_subplot(111)
        ax.bar(np.arange(1,len(lamb)+1), lamb)
        ax.plot(range(1,len(l1)+1),l1)
        ax.plot(range(1,len(l2)+1),l2)
        ax.set_xlabel('Eigenvalue index',fontsize=15)
        ax.set_ylabel('Eigenvalues',fontsize=15)
        ax.set_title('L Method implies data is %d dimensional' %(true_dim+1),fontsize=15)
        plt.xticks(np.arange(1, len(lamb)+1, step=1))
        plt.tight_layout()
        plt.show()
        plt.close()

    return true_dim+1

####################################################################################################################
# 4.5: Nystrom Extension
def nystrom(ts,tau,dim,lamb,psi,dvecs,train,eps,save_loc='N',save=True,plots=False,plot_stride=100):

    """
    Compute the Nystrom extension of the diffusion map.

    Args:
        ts: The time series name (string).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        lamb: The eigenvalues of the diffusion map (array).
        psi: The eigenvectors of the diffusion map (array).
        dvecs: The delay vectors of the time series (array).
        train: The training data (array).
        eps: The optimal epsilon value for the diffusion map (float).
        save_loc: The location to save the Nystrom extension (string).
        save: A boolean to save the Nystrom extension (boolean).
        plots: A boolean to plot the Nystrom extension (boolean).
        plot_stride: The stride of the Nystrom extension plot (int).

    Returns:
        array: The Nystrom extension of the diffusion map vectors.
    """

    dvecs_ss = dvecs[train]
    z = []
    for i in range(dvecs.shape[0]):
        Euc_i = euclidean_distances(dvecs_ss,dvecs[i].reshape(1,-1)) 
        Euc_i = np.exp(-Euc_i**2/(2*eps))
        Euc_i /= np.sum(Euc_i)
        psi_Nystrom = np.divide( np.matmul(Euc_i.T,psi), lamb)
        z.append(psi_Nystrom)
    z = np.array(z)
    z = z.reshape(z.shape[0],-1)
    
    # saving dMap embedding of all data
    np.savez('%s/data_%s_tau%d_dim%d.npz' %(save_loc,ts,tau,dim), z=z)

    if plots==True:

        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        i,j,k=1,2,3
        ax.scatter(psi[::plot_stride,i], psi[::plot_stride,j], psi[::plot_stride,k])
        ax.set_xlabel(r'$\Psi_%s$' %str(i),fontsize=15)
        ax.set_ylabel(r'$\Psi_%s$' %str(j),fontsize=15)
        ax.set_zlabel(r'$\Psi_%s$' %str(k),fontsize=15)
        ax.set_title('Nystrom Extension of Embedding',fontsize=15)
        plt.tight_layout()
        plt.draw()
        plt.show()
        plt.close()
        plt.savefig('%s/nystrom_%s_tau%d_dim%d.pdf' %(save_loc,ts,tau,dim))
    
    return z
####################################################################################################################
#Loading Functions
# 5.1: Loading Manifolds
def load_manif(load_loc,n_features,first_feature=0,z='empty'):
    """
    Load a manifold from an npz file and return the features we wish to use.

    Args:
        load_loc: location of the npz file containing the manifold (string).
        n_features: The number of features in the manifold. This should correspond the the results of the impplied dimensionality from the L method (int).
        first_feature: The index of the first feature we wish to use (int).
        z: The name of the key in the npz file that contains the manifold. If 'empty' then the first key will be used (string).

    Returns:
        array: The manifold with the desired features
    """
    npzfile = np.load(load_loc)
    if z=='empty':
        key = list(npzfile.keys())[0]
    else:
        key = z
    z_aa = npzfile[key]
    z_aa_features_idx = np.array(range(first_feature,n_features+first_feature))
    z_aa = z_aa[:,z_aa_features_idx]
    return z_aa

####################################################################################################################
# 6.1: Nonlinear Mapping
def train_and_evaluate_model(Z_out, load_loc,idim,odim, hidden_layers=[50,50,50,50,50], epochs=250, learning_rate=0.01, first_feature=1, save=False, save_loc='N', split=0.5):

    """
    Train a neural network to predict the next state of a dynamical system.

    Args:
        Z_out: The output of the diffusion map (array).
        load_loc: The location of the delay vectors (string).
        idim: The input dimension of the neural network (int).
        odim: The output dimension of the neural network (int).
        hidden_layers: Hidden layers in the neural network (list).
        epochs: The number of epochs to train the neural network (int).
        learning_rate: The learning rate of the neural network (float).
        first_feature: The index of the first feature we wish to use (int).
        save: A boolean to save the neural network output (boolean).
        save_loc: The location to save the neural network output (string).
        split: The proportion of data to use for training (float).


    Returns:
        float: The loss of the neural network.
        float: The root mean squared error of the neural network.
    """ 
    
    Y = Z_out
    # read in Z_in
    X = load_manif(load_loc, idim, first_feature=first_feature)
    Y = Y[int(Y.shape[0] - X.shape[0]):, 0:odim] ### was 0:2
    # training only on first split of data
    Y_train = Y[0:int(len(Y) * split)]
    X_train = X[0:int(len(X) * split)]
    Y_test = Y[int(len(Y) * split):]
    X_test = X[int(len(X) * split):]
    
    # scale and transform data
    x_scaler, y_scaler = MinMaxScaler(), MinMaxScaler()
    x_train = x_scaler.fit_transform(X_train)
    y_train = y_scaler.fit_transform(Y_train)

    # convert to PyTorch
    x_train = torch.from_numpy(x_train).float()
    y_train = torch.from_numpy(y_train).float()
    input_dim = idim
    output_dim = odim
    print('input dim is: ', input_dim), print('output dim is: ', output_dim), print('hidden layers are: ', hidden_layers)
    class Net(nn.Module):
        def __init__(self, input_dim, hidden_layers, output_dim):
            super(Net, self).__init__()
            layers = []
            in_dim = input_dim
            for h_dim in hidden_layers:
                layers.append(nn.Linear(in_dim, h_dim))
                layers.append(nn.ReLU())
                in_dim = h_dim
            layers.append(nn.Linear(in_dim, output_dim))
            self.network = nn.Sequential(*layers)

        def forward(self, x):
            return self.network(x)

    # make network
    net = Net(input_dim, hidden_layers, output_dim)

    # define loss function and optimizer
    criterion = nn.MSELoss()
    optimizer = Adam(net.parameters(), lr=learning_rate)

    # build the model and train
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = net(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, loss {loss.item()}')
    print(f'Finished epoch {epoch}, latest loss {loss}')

    # scale and transform test data
    x_test = x_scaler.transform(X_test)
    y_test = y_scaler.transform(Y_test)

    # convert test to torch
    x_test = torch.from_numpy(x_test).float()
    y_test = torch.from_numpy(y_test).float()

    # pass test through network
    outputs_test = net(x_test)
    X_SCL = outputs_test.detach().numpy()
    X_USC = y_scaler.inverse_transform(X_SCL)
    if save:
        np.savez(save_loc, z_out=outputs_test.detach().numpy())
    # compute loss
    LOSS = criterion(outputs_test, y_test).item()
    RMSE = np.mean(np.sqrt((X_USC - Y_test) ** 2))

    return LOSS, RMSE
    
def train_and_evaluate_model2(Z_out, load_loc,idim,odim,trial, hidden_layers=[50,50,50,50,50], epochs=250, learning_rate=0.01, first_feature=0, save=False, save_loc='N', divisions=3):

    """
    Train a neural network to predict the next state of a dynamical system.

    Args:
        Z_out: The output of the diffusion map (array).
        load_loc: The location of the delay vectors (string).
        idim: The input dimension of the neural network (int).
        odim: The output dimension of the neural network (int).
        hidden_layers: Hidden layers in the neural network (list).
        epochs: The number of epochs to train the neural network (int).
        learning_rate: The learning rate of the neural network (float).
        first_feature: The index of the first feature we wish to use (int).
        save: A boolean to save the neural network output (boolean).
        save_loc: The location to save the neural network output (string).
        split: The proportion of data to use for training (float).


    Returns:
        float: The loss of the neural network.
        float: The root mean squared error of the neural network.
    """ 
    
    Y = Z_out
    # read in Z_in
    X = load_manif(load_loc, idim, first_feature=first_feature)

    Y = Y[int(Y.shape[0] - X.shape[0]):, 0:odim] ### was 0:2
    print(Y.shape, X.shape)
    lengths = len(Y)
    samples  =int(lengths/divisions)
    print(lengths, samples)
    train_indices = np.arange(int(trial*samples),int((trial+1)*samples))
    # training only on first split of data
    Y_train = Y[train_indices,:]#.reshape(-1,4)
    X_train = X[train_indices,:]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape)
    # make the test the rest of our data points
    #mask = np.isin(Y, Y[int(trial*samples):int((trial+1)*samples),:], invert=True)
    mask = np.ones(len(Y), dtype=bool)  # Start with all True
    mask[train_indices] = False         # Set train indices to False
    Y_test = Y[mask]#.reshape(-1,4)
    #mask = np.isin(X, X[int(trial*samples):int((trial+1)*samples),:], invert=True)
    X_test = X[mask]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape, X_test.shape, Y_test.shape)

    # scale and transform data
    x_scaler, y_scaler = MinMaxScaler(), MinMaxScaler()
    x_train = x_scaler.fit_transform(X_train)
    y_train = y_scaler.fit_transform(Y_train)

    # convert to PyTorch
    x_train = torch.from_numpy(x_train).float()
    y_train = torch.from_numpy(y_train).float()
    input_dim = idim
    output_dim = odim
    print('input dim is: ', input_dim), print('output dim is: ', output_dim), print('hidden layers are: ', hidden_layers)
    class Net(nn.Module):
        def __init__(self, input_dim, hidden_layers, output_dim):
            super(Net, self).__init__()
            layers = []
            in_dim = input_dim
            for h_dim in hidden_layers:
                layers.append(nn.Linear(in_dim, h_dim))
                layers.append(nn.ReLU())
                in_dim = h_dim
            layers.append(nn.Linear(in_dim, output_dim))
            self.network = nn.Sequential(*layers)

        def forward(self, x):
            return self.network(x)

    # make network
    net = Net(input_dim, hidden_layers, output_dim)

    # define loss function and optimizer
    criterion = nn.MSELoss()
    optimizer = Adam(net.parameters(), lr=learning_rate)

    # build the model and train
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = net(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, loss {loss.item()}')
    print(f'Finished epoch {epoch}, latest loss {loss}')

    # scale and transform test data
    x_test = x_scaler.transform(X_test)
    y_test = y_scaler.transform(Y_test)

    # convert test to torch
    x_test = torch.from_numpy(x_test).float()
    y_test = torch.from_numpy(y_test).float()

    # pass test through network
    outputs_test = net(x_test)
    X_SCL = outputs_test.detach().numpy()
    X_USC = y_scaler.inverse_transform(X_SCL)
    if save:
        np.savez(save_loc, z_out=outputs_test.detach().numpy())
    # compute loss
    LOSS = criterion(outputs_test, y_test).item()
    RMSE = np.mean(np.sqrt((X_USC - Y_test) ** 2))

    return LOSS, RMSE

def train_and_evaluate_model3(Z_out, Z_in,idim,odim,trial, hidden_layers=[50,50,50,50,50], epochs=250, learning_rate=0.01, first_feature=0, save=False, save_loc='N', divisions=3):

    """
    Train a neural network to predict the next state of a dynamical system.

    Args:
        Z_out: The output of the diffusion map (array).
        load_loc: The location of the delay vectors (string).
        idim: The input dimension of the neural network (int).
        odim: The output dimension of the neural network (int).
        hidden_layers: Hidden layers in the neural network (list).
        epochs: The number of epochs to train the neural network (int).
        learning_rate: The learning rate of the neural network (float).
        first_feature: The index of the first feature we wish to use (int).
        save: A boolean to save the neural network output (boolean).
        save_loc: The location to save the neural network output (string).
        split: The proportion of data to use for training (float).


    Returns:
        float: The loss of the neural network.
        float: The root mean squared error of the neural network.
    """ 
    
    Y = Z_out
    # read in Z_in
    X = Z_in

    Y = Y[int(Y.shape[0] - X.shape[0]):, 0:odim] ### was 0:2
    print(Y.shape, X.shape)
    lengths = len(Y)
    samples  =int(lengths/divisions)
    print(lengths, samples)
    train_indices = np.arange(int(trial*samples),int((trial+1)*samples))
    # training only on first split of data
    Y_train = Y[train_indices,:]#.reshape(-1,4)
    X_train = X[train_indices,:]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape)
    # make the test the rest of our data points
    #mask = np.isin(Y, Y[int(trial*samples):int((trial+1)*samples),:], invert=True)
    mask = np.ones(len(Y), dtype=bool)  # Start with all True
    mask[train_indices] = False         # Set train indices to False
    Y_test = Y[mask]#.reshape(-1,4)
    #mask = np.isin(X, X[int(trial*samples):int((trial+1)*samples),:], invert=True)
    X_test = X[mask]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape, X_test.shape, Y_test.shape)

    # scale and transform data
    x_scaler, y_scaler = MinMaxScaler(), MinMaxScaler()
    x_train = x_scaler.fit_transform(X_train)
    y_train = y_scaler.fit_transform(Y_train)

    # convert to PyTorch
    x_train = torch.from_numpy(x_train).float()
    y_train = torch.from_numpy(y_train).float()
    input_dim = idim
    output_dim = odim
    print('input dim is: ', input_dim), print('output dim is: ', output_dim), print('hidden layers are: ', hidden_layers)
    class Net(nn.Module):
        def __init__(self, input_dim, hidden_layers, output_dim):
            super(Net, self).__init__()
            layers = []
            in_dim = input_dim
            for h_dim in hidden_layers:
                layers.append(nn.Linear(in_dim, h_dim))
                layers.append(nn.ReLU())
                in_dim = h_dim
            layers.append(nn.Linear(in_dim, output_dim))
            self.network = nn.Sequential(*layers)

        def forward(self, x):
            return self.network(x)

    # make network
    net = Net(input_dim, hidden_layers, output_dim)

    # define loss function and optimizer
    criterion = nn.MSELoss()
    optimizer = Adam(net.parameters(), lr=learning_rate)

    # build the model and train
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = net(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, loss {loss.item()}')
    print(f'Finished epoch {epoch}, latest loss {loss}')

    # scale and transform test data
    x_test = x_scaler.transform(X_test)
    y_test = y_scaler.transform(Y_test)

    # convert test to torch
    x_test = torch.from_numpy(x_test).float()
    y_test = torch.from_numpy(y_test).float()

    # pass test through network
    outputs_test = net(x_test)
    X_SCL = outputs_test.detach().numpy()
    X_USC = y_scaler.inverse_transform(X_SCL)
    if save:
        np.savez(save_loc, z_out=outputs_test.detach().numpy())
    # compute loss
    LOSS = criterion(outputs_test, y_test).item()
    RMSE = np.mean(np.sqrt((X_USC - Y_test) ** 2))

    return LOSS, RMSE

def train_and_evaluate_model4(Z_out, load_loc,idim,odim,trial, hidden_layers=[50,50,50,50,50], epochs=250, learning_rate=0.01, first_feature=0, save=False, save_loc='N', divisions=3):

    
    """
    Train a neural network to predict the next state of a dynamical system.

    Args:
        Z_out: The output of the diffusion map (array).
        load_loc: The location of the delay vectors (string).
        idim: The input dimension of the neural network (int).
        odim: The output dimension of the neural network (int).
        hidden_layers: Hidden layers in the neural network (list).
        epochs: The number of epochs to train the neural network (int).
        learning_rate: The learning rate of the neural network (float).
        first_feature: The index of the first feature we wish to use (int).
        save: A boolean to save the neural network output (boolean).
        save_loc: The location to save the neural network output (string).
        split: The proportion of data to use for training (float).


    Returns:
        float: The loss of the neural network.
        float: The root mean squared error of the neural network.
    """ 
    
    Y = Z_out
    # read in Z_in
    X = load_manif(load_loc, idim, first_feature=first_feature)

    Y = Y[int(Y.shape[0] - X.shape[0]):, 0:odim] ### was 0:2
    print(Y.shape, X.shape)
    lengths = len(Y)
    samples  =int(lengths/divisions)
    print(lengths, samples)
    train_indices = np.arange(int(trial*samples),int((trial+1)*samples))
    # training only on first split of data
    Y_train = Y[train_indices,:]#.reshape(-1,4)
    X_train = X[train_indices,:]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape)
    # make the test the rest of our data points
    #mask = np.isin(Y, Y[int(trial*samples):int((trial+1)*samples),:], invert=True)
    mask = np.ones(len(Y), dtype=bool)  # Start with all True
    mask[train_indices] = False         # Set train indices to False
    Y_test = Y[mask]#.reshape(-1,4)
    #mask = np.isin(X, X[int(trial*samples):int((trial+1)*samples),:], invert=True)
    X_test = X[mask]#.reshape(-1,4)
    print(X_train.shape, Y_train.shape, X_test.shape, Y_test.shape)

    # scale and transform data
    x_scaler, y_scaler = MinMaxScaler(), MinMaxScaler()
    x_train = x_scaler.fit_transform(X_train)
    y_train = y_scaler.fit_transform(Y_train)

    # convert to PyTorch
    x_train = torch.from_numpy(x_train).float()
    y_train = torch.from_numpy(y_train).float()
    input_dim = idim
    output_dim = odim
    print('input dim is: ', input_dim), print('output dim is: ', output_dim), print('hidden layers are: ', hidden_layers)
    class Net(nn.Module):
        def __init__(self, input_dim, hidden_layers, output_dim):
            super(Net, self).__init__()
            layers = []
            in_dim = input_dim
            for h_dim in hidden_layers:
                layers.append(nn.Linear(in_dim, h_dim))
                layers.append(nn.Tanh())
                in_dim = h_dim
            layers.append(nn.Linear(in_dim, output_dim))
            self.network = nn.Sequential(*layers)

        def forward(self, x):
            return self.network(x)

    # make network
    net = Net(input_dim, hidden_layers, output_dim)

    # define loss function and optimizer
    criterion = nn.MSELoss()
    optimizer = Adam(net.parameters(), lr=learning_rate)

    # build the model and train
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = net(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        if epoch % 100 == 0:
            print(f'Epoch {epoch}, loss {loss.item()}')
    print(f'Finished epoch {epoch}, latest loss {loss}')

    # scale and transform test data
    x_test = x_scaler.transform(X_test)
    y_test = y_scaler.transform(Y_test)

    # convert test to torch
    x_test = torch.from_numpy(x_test).float()
    y_test = torch.from_numpy(y_test).float()

    # pass test through network
    outputs_test = net(x_test)
    X_SCL = outputs_test.detach().numpy()
    X_USC = y_scaler.inverse_transform(X_SCL)
    if save:
        np.savez(save_loc, z_out=outputs_test.detach().numpy())
    # compute loss
    LOSS = criterion(outputs_test, y_test).item()
    RMSE = np.mean(np.sqrt((X_USC - Y_test) ** 2))

    return LOSS, RMSE

####################################################################################################################
# 7.1 3D Probability maps 
  
def pmap3D(z,bins=100,mapdims=[1,2,3],tau=100,dim=7,crit=0): 
    """
    Compute the probability map of a manifold.

    Args:
        z: The manifold (array).
        bins: The number of bins in the probability map (int).
        mapdims: The dimensions of the manifold to use in the probability map (list).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        crit: The critical cutoff value for the probability map used to cutoff insignificant states (float).

    Returns:      
        fedfxy_c: The log of the probability of occupying any given xy bin (array).
        fedfyz_c: The log of the probability of occupying any given yz bin (array).
        locx: The x bins locations of the manifold in the probability map (array).
        locy: The y bins locations of the manifold in the probability map (array).
        locz: The z bins locations of the manifold in the probability map (array).
    """ 

    #decide bounds
    z1min,z1max,z2min,z2max,z3min,z3max = np.nanmin(z[:,mapdims[0]]),np.nanmax(z[:,mapdims[0]]),np.nanmin(z[:,mapdims[1]]),np.nanmax(z[:,mapdims[1]]),np.nanmin(z[:,mapdims[2]]),np.nanmax(z[:,mapdims[2]])
    
    # decide edges
    xedges = np.linspace(z1min,z1max,bins)
    yedges = np.linspace(z2min,z2max,bins)
    zedges = np.linspace(z3min,z3max,bins)

    fexy,xxedges,xyedges = np.histogram2d(z[:,mapdims[0]],z[:,mapdims[1]],range=[[z1min,z1max],[z2min,z2max]],bins=(xedges,yedges))
    feyz,zyedges,zzedges = np.histogram2d(z[:,mapdims[1]],z[:,mapdims[2]],range=[[z2min,z2max],[z3min,z3max]],bins=(yedges,zedges))
    fexy = fexy/np.sum(fexy)
    feyz = feyz/np.sum(feyz)
    #if we choose a critical cutoff 'crit'
    if crit!=0:
        critmatxy = np.ones([len(fexy)+2,len(fexy)+2])
        critmatyz = np.ones([len(fexy)+2,len(fexy)+2])
        #iterate over all bins
        for jfk in range(len(fexy)):
            for kfj in range(len(fexy)):
                #set extremely rarely sampled states to zero
                if fexy[jfk,kfj]<=crit:
                   fexy[jfk,kfj] = 0
                   critmatxy[jfk,kfj] = 0
                if feyz[jfk,kfj]<=crit:
                   feyz[jfk,kfj] = 0
                   critmatyz[jfk,kfj] = 0

        
    fexy = fexy/np.sum(fexy)
    fedfxy= -np.log(fexy.T)
    #make inversion
    fedfxy_c = np.flip(deepcopy(fedfxy), axis=0)

    
    feyz = feyz/np.sum(feyz)
    fedfyz= -np.log(feyz.T)
    #make inversion
    fedfyz_c = np.flip(deepcopy(fedfyz), axis=0)


    zpx,zpy,zpz = [],[],[]
    for i in range(len(z)):
        zpx.append(z[i][mapdims[0]])
        zpy.append(z[i][mapdims[1]])
        zpz.append(z[i][mapdims[2]])   
 

    locx,locy,locz = [],[],[]
    for i in range(len(z)):
        locx.append((zpx[i]-z1min)/(z1max-z1min)*bins)
        locy.append((zpy[i]-z2min)/(z2max-z2min)*bins)
        locz.append((zpz[i]-z3min)/(z3max-z3min)*bins)

    #eliminate rare sampled states from locxyz by multiplying by critmat
    if crit!=0:
        for i in range(len(locx)):
            cordsxy = [int(np.floor(locx[i])),int(np.floor(locy[i]))]
            cordsyz = [int(np.floor(locy[i])),int(np.floor(locz[i]))]

            if critmatxy[cordsxy[0],cordsxy[1]]!=1:
                locx[i]=np.nan
                locy[i]=np.nan
                locz[i]=np.nan
            if critmatyz[cordsyz[0],cordsyz[1]]!=1:
                locx[i]=np.nan
                locy[i]=np.nan
                locz[i]=np.nan

    return fedfxy_c,fedfyz_c,locx,locy,locz

####################################################################################################################
# 7.2: 2D Probability maps

def pmap2D(z,bins=100,mapdims=[1,2],tau=100,dim=7,crit=0): 
    """
    Compute the probability map of a manifold for 2D visualisation.

    Args:
        z: The manifold (array).
        bins: The number of bins in the probability map (int).
        mapdims: The dimensions of the manifold to use in the probability map (list).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        crit: The critical cutoff value for the probability map (float).

    Returns:    
        fedfxy_c: The log of the probability of occupying any given bin (array).
        locx: The x bins locations of the manifold in the probability map (array).
        locy: The y bins locations of the manifold in the probability map (array).
    """ 

    #decide bounds
    z1min,z1max,z2min,z2max = np.nanmin(z[:,mapdims[0]]),np.nanmax(z[:,mapdims[0]]),np.nanmin(z[:,mapdims[1]]),np.nanmax(z[:,mapdims[1]])
    
    # decide edges
    xedges = np.linspace(z1min,z1max,bins)
    yedges = np.linspace(z2min,z2max,bins)

    fexy,xxedges,xyedges = np.histogram2d(z[:,mapdims[0]],z[:,mapdims[1]],range=[[z1min,z1max],[z2min,z2max]],bins=(xedges,yedges))

    fexy = fexy/np.sum(fexy)
    #if we choose a critical cutoff 'crit'
    if crit!=0:
        critmatxy = np.ones([len(fexy)+2,len(fexy)+2])
        #iterate over all bins
        for jfk in range(len(fexy)):
            for kfj in range(len(fexy)):
                #set extremely rarely sampled states to zero
                if fexy[jfk,kfj]<=crit:
                   fexy[jfk,kfj] = 0
                   critmatxy[jfk,kfj] = 0

        
    fexy = fexy/np.sum(fexy)
    fedfxy= -np.log(fexy.T)
    #make inversion
    fedfxy_c = np.flip(deepcopy(fedfxy), axis=0)


    zpx,zpy = [],[]
    for i in range(len(z)):
        zpx.append(z[i][mapdims[0]])
        zpy.append(z[i][mapdims[1]])


    locx,locy = [],[]
    for i in range(len(z)):
        locx.append((zpx[i]-z1min)/(z1max-z1min)*bins)
        locy.append((zpy[i]-z2min)/(z2max-z2min)*bins)
        #eliminate rare sampled states from locxyz by multiplying by critmat
    if crit!=0:
        for i in range(len(locx)):
            cordsxy = [int(np.floor(locx[i])),int(np.floor(locy[i]))]
            if critmatxy[cordsxy[0],cordsxy[1]]!=1:
                locx[i]=np.nan
                locy[i]=np.nan


    return fedfxy_c,locx,locy

####################################################################################################################
# 7.3: Probability maps for visualisation

def pmap_vis(z, ts, bins=100, mapdims=[1, 2], tau=100, dim=7, save_loc='N', save=False, plots=False, axes=[r'$\Psi_1$', r'$\Psi_2$']):
    """
    Compute the probability map of a manifold for visualisation.

    Args:
        z: The manifold (array).
        ts: The time series name (string).
        bins: The number of bins in the probability map (int).
        mapdims: The dimensions of the manifold to use in the probability map (list).
        tau: The time delay of the time series (int).
        dim: The dimensionality of the time series (int).
        save_loc: The location to save the probability map (string).
        save: A boolean to save the probability map (boolean).
        plots: A boolean to plot the probability map (boolean).
        axes: The axes of the probability map (list).
    
    Returns:
        fe: The probability of occupying any given bin (array).
        feDf: The log of the probability of occupying any given bin (array).
    """
    # decide bounds
    z1min, z1max, z2min, z2max = np.nanmin(z[:, mapdims[0]]), np.nanmax(z[:, mapdims[0]]), np.nanmin(z[:, mapdims[1]]), np.nanmax(z[:, mapdims[1]])
    # decide edges
    xedges = np.linspace(z1min, z1max, bins)
    yedges = np.linspace(z2min, z2max, bins)
    fe, xedges, yedges = np.histogram2d(z[:, mapdims[0]], z[:, mapdims[1]], range=[[z1min, z1max], [z2min, z2max]], bins=(xedges, yedges))
    fe = fe / np.sum(fe)
    fedf = -np.log(fe.T)
    # make inversion
    fedfc = np.flip(deepcopy(fedf), axis=0)
    # plot
    fig, ax = plt.subplots()
    im = ax.imshow(fedfc, interpolation='nearest', cmap='viridis')  # , origin='low', extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]])
    ax.set_xlabel(axes[0])
    ax.set_ylabel(axes[1])
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = fig.colorbar(im, cax=cax)
    cbar.set_label('-log(P('+axes[0]+','+axes[1]+'))')
    if save:
        ax.figure.savefig('%s/PMAP_%s_tau%d_dim%d.pdf' % (save_loc, ts, tau, dim))
        np.savez('%s/PMAP_%s_tau%d_dim%d.npz' % (save_loc, ts, tau, dim), fe=fe, fedf=fedf)
        
    return fe, fedf





####################################################################################################################
# Bootstrap & evaluation helpers
def block_bootstrap_mad(diff, mask=None, n_bootstrap=1000, block_size=5000, seed=0):
    """Block bootstrap of the mean absolute deviation on a per-frame difference array.

    Produces 95% confidence intervals for MAD = mean(|diff|) by resampling
    contiguous BLOCKS of frames with replacement (preserves time autocorrelation).

    Args:
        diff (array): Per-frame difference, shape (n_frames, n_atoms, 3) OR
            (n_frames, n_atoms*3). Pre-aligned to the test set of interest.
        mask (array of bool): Optional atom-selection mask of length n_atoms.
            If None, evaluates over all atoms.
        n_bootstrap (int): Number of bootstrap resamples. Default 200.
        block_size (int): Block length in frames. Choose >= autocorrelation
            time of the trajectory. Default 5000.
        seed (int): RNG seed for reproducibility. Default 0.

    Returns:
        (mean MAD, [ci_lo_2.5, ci_hi_97.5]) where mean is the bootstrap mean
        and the CI is the central 95% percentile interval.
    """
    if diff.ndim == 2:
        # reshape (n_frames, n_atoms*3) into (n_frames, n_atoms, 3)
        n_atoms = diff.shape[1] // 3
        diff = diff.reshape(diff.shape[0], n_atoms, 3)
    n_test = diff.shape[0]
    n_atoms = diff.shape[1]
    if mask is None:
        mask = np.ones(n_atoms, dtype=bool)
    n_blocks = n_test // block_size
    if n_blocks < 2:
        raise ValueError("trajectory too short for chosen block_size")
    rng = np.random.default_rng(seed)
    block_starts = np.arange(n_blocks) * block_size
    boots = np.empty(n_bootstrap)
    for b in range(n_bootstrap):
        sel = rng.choice(block_starts, size=n_blocks, replace=True)
        sample = np.concatenate([diff[s:s + block_size, mask, :] for s in sel],
                                  axis=0)
        boots[b] = float(np.abs(sample).mean())
    return float(boots.mean()), tuple(np.percentile(boots, [2.5, 97.5]))


def per_component_mad(recon, true, masks):
    """Compute MAD per atom-group on a recon vs true array pair.

    Args:
        recon, true (array): shape (n_frames, n_atoms, 3) - recon and true must
            be already aligned to the same window.
        masks (dict[str, array]): atom-group name -> bool mask of length n_atoms.

    Returns:
        dict[str, float]: MAD per group.
    """
    out = {}
    for name, mask in masks.items():
        d = recon[:, mask, :] - true[:, mask, :]
        out[name] = float(np.abs(d).mean())
    return out


def parse_pdb_heavy(pdb_path):
    """Parse heavy-atom (non-hydrogen) names + elements from a PDB file.

    Args:
        pdb_path (str): path to a PDB file.

    Returns:
        names (np.ndarray of str): heavy-atom names (e.g. "CA", "CB", ...)
            in PDB order, hydrogens excluded.
        elements (np.ndarray of str): corresponding element symbols.
    """
    names, elems = [], []
    for L in open(pdb_path):
        if L.startswith("ATOM"):
            names.append(L[12:16].strip())
            e = L[76:78].strip() if len(L) >= 78 else ""
            if not e:
                e = (L[12:16].strip().lstrip("0123456789") or "X")[0]
            elems.append(e)
    names = np.array(names); elems = np.array(elems)
    keep = (elems != "H")
    return names[keep], elems[keep]


def block_bootstrap_rmsd(diff, mask=None, n_bootstrap=1000, block_size=5000, seed=0):
    """Block-bootstrap 95% CI on RMSD = sqrt(mean(diff^2)).

    diff : (n_frames, n_atoms, 3) or (n_frames, n_atoms*3) per-frame difference.
    mask : optional bool array of length n_atoms (which atoms to score).
    Returns (mean RMSD, (ci_lo, ci_hi)).
    """
    if diff.ndim == 2:
        n_atoms = diff.shape[1] // 3
        diff = diff.reshape(diff.shape[0], n_atoms, 3)
    n_atoms = diff.shape[1]
    if mask is None:
        mask = np.ones(n_atoms, dtype=bool)
    n_test = diff.shape[0]
    n_blocks = n_test // block_size
    if n_blocks < 2:
        raise ValueError("trajectory too short for chosen block_size")
    rng = np.random.default_rng(seed)
    starts = np.arange(n_blocks) * block_size
    boots = np.empty(n_bootstrap)
    for b in range(n_bootstrap):
        sel = rng.choice(starts, size=n_blocks, replace=True)
        sample = np.concatenate([diff[s:s + block_size, mask, :] for s in sel],
                                  axis=0)
        boots[b] = float(np.sqrt((sample ** 2).mean()))
    return float(boots.mean()), tuple(np.percentile(boots, [2.5, 97.5]))


def lift_returns_ann(psi, returns, hidden=(50, 50, 50, 50, 50), epochs=250,
                       lr=0.01, batch_size=512, train_frac=0.5, seed=0):
    """Train a small ANN to lift from a manifold projection psi to scalar returns.

    psi      : (n, d_psi) manifold coordinates (e.g. leading dMap eigvecs).
    returns  : (n,) scalar returns aligned with psi.
    Returns predicted_returns of shape (n,) covering the entire input.
    """
    import torch
    import torch.nn as nn
    from sklearn.preprocessing import MinMaxScaler

    torch.manual_seed(seed); np.random.seed(seed)
    n = psi.shape[0]; n_train = int(n * train_frac)
    x_sc = MinMaxScaler(); y_sc = MinMaxScaler()
    Xn = x_sc.fit_transform(psi[:n_train])
    Yn = y_sc.fit_transform(returns[:n_train].reshape(-1, 1))
    Xfull = x_sc.transform(psi)

    layers = []
    prev = psi.shape[1]
    for h in hidden:
        layers += [nn.Linear(prev, h), nn.ReLU()]
        prev = h
    layers += [nn.Linear(prev, 1)]
    model = nn.Sequential(*layers)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    Xt = torch.from_numpy(Xn).float()
    Yt = torch.from_numpy(Yn).float()
    for ep in range(epochs):
        perm = torch.randperm(Xt.shape[0])
        for i in range(0, Xt.shape[0], batch_size):
            idx = perm[i:i+batch_size]
            opt.zero_grad()
            loss = ((model(Xt[idx]) - Yt[idx]) ** 2).mean()
            loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        pred = model(torch.from_numpy(Xfull).float()).numpy()
    return y_sc.inverse_transform(pred).flatten()


def delta_deformation(canonical_xyz, true_xyz, recon_xyz, scale=1.0):
    """Apply per-atom delta from (recon - true) as a perturbation to a canonical
    reference structure.

    canonical_xyz : (n_atoms, 3) reference coords (e.g. canonical Villin folded).
    true_xyz      : (n_atoms, 3) MD-truth coords for this frame.
    recon_xyz     : (n_atoms, 3) reconstruction coords for this frame.
    scale         : optional damping factor in [0, 1] so the deformation stays
                    visible without breaking secondary-structure assignment.
    Returns       : (n_atoms, 3) deformed coords.
    """
    return canonical_xyz + scale * (recon_xyz - true_xyz)


def kabsch_align(P, Q):
    """Kabsch superposition: rotate + translate P onto Q.

    P, Q : (n_atoms, 3) coordinate arrays.
    Returns the rotated+translated P such that ((kabsch_align(P, Q) - Q)**2).sum()
    is minimised over rigid transforms.

    For a structural RMSD, *always* superpose with this function before
    differencing — otherwise the reported RMSD reflects rigid-body drift,
    not the deformation of interest.
    """
    Pc = P - P.mean(0)
    Qc = Q - Q.mean(0)
    H = Pc.T @ Qc
    U, _, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(Vt.T @ U.T))
    R = Vt.T @ np.diag([1.0, 1.0, d]) @ U.T
    return Pc @ R.T + Q.mean(0)


def _kabsch_chunk(recon_chunk, true_chunk, mask):
    """Internal: vectorised Kabsch-aligned diff over a single chunk."""
    if mask is None:
        rec_cent = recon_chunk.mean(axis=1, keepdims=True)
        tru_cent = true_chunk.mean(axis=1, keepdims=True)
        Rc = recon_chunk - rec_cent
        Tc = true_chunk  - tru_cent
        H = np.einsum('fai,faj->fij', Rc, Tc)
    else:
        rec_cent = recon_chunk[:, mask, :].mean(axis=1, keepdims=True)
        tru_cent = true_chunk[:,  mask, :].mean(axis=1, keepdims=True)
        Rc = recon_chunk - rec_cent
        Tc = true_chunk  - tru_cent
        H = np.einsum('fai,faj->fij', Rc[:, mask, :], Tc[:, mask, :])
    U, _, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(np.einsum('fij,fkj->fik', Vt, U)))
    D = np.zeros((H.shape[0], 3, 3))
    D[:, 0, 0] = 1.0; D[:, 1, 1] = 1.0; D[:, 2, 2] = d
    R = np.einsum('fji,fjk,flk->fil', Vt, D, U)
    Rc_rotated = np.einsum('fai,fji->faj', Rc, R)
    return (Rc_rotated + tru_cent) - true_chunk


def kabsch_aligned_diff(recon, true, mask=None, chunk=20000, verbose=False):
    """Per-frame Kabsch-superposed (recon - true) over an optional atom subset.

    recon, true : (n_frames, n_atoms, 3) trajectories.
    mask        : optional bool mask of length n_atoms used for the
                  superposition (e.g. backbone-only). Differences are
                  returned for ALL atoms after the alignment; pass mask=None
                  to use all atoms.
    chunk       : number of frames per batch. Defaults to 20,000 to bound
                  peak memory at ~1.5 GB for n_atoms=287; raise on a high-
                  memory node, lower if you OOM.
    verbose     : if True, print one line per chunk so a user with a
                  300k-frame trajectory sees progress.
    Returns     : (n_frames, n_atoms, 3) post-alignment differences.

    Implementation: vectorised batched SVD over the frame axis, processed
    in chunks. The non-chunked equivalent matches scalar Kabsch to ~1e-15.
    """
    n = recon.shape[0]
    out = np.empty_like(recon)
    n_chunks = (n + chunk - 1) // chunk
    for i, s in enumerate(range(0, n, chunk)):
        e = min(s + chunk, n)
        out[s:e] = _kabsch_chunk(recon[s:e], true[s:e], mask)
        if verbose:
            print(f'  kabsch chunk {i+1}/{n_chunks}: frames {s}-{e}', flush=True)
    return out