library("Langevin")
# Load library used for plotting
library("plotrix")

# Set the number of bins to estimate drift and diffusion on
bins <- 40
# Set the time lags for computing the conditional moments
steps <- c(1:3)
# Set the sampling frequency
sf <- 1000

##########################################
# One dimensional example
##########################################

# Generate the time series
set.seed(4711)
Ux <- timeseries1D(N = 1e7, d11 = 1, d13 = -1, d22 = 1, d20 = 1, sf = sf)

# Plot time series and probability density function (PDF)
op <- par(no.readonly = TRUE)
par(mar = c(4.2, 5.4, 0.5, 0.5))
layout(matrix(c(1, 1, 2), 1, 3))
plot((1:1e5)/sf, Ux[1:1e5], xlab = "t [s]", ylab = "x [a.u.]", t = 'l')
plot(density(Ux), xlab = "x [a.u.]", ylab = "density", t = 'l',main = "")
par(op)

# Esitmate drift and diffusion coefficients
ests <- Langevin1D(Ux, bins, steps)
attach(ests)

# Plot drift and diffusion coefficients with errors and theoretical expecations
# for comparison
op <- par(no.readonly = TRUE)
par(mar = c(4.2, 5.4, 0.5, 0.5))
par(mfrow = c(1, 2))
plotCI(mean_bin, D1, uiw = eD1, xlab = "x [a.u.]",
       ylab = expression(paste("Drift coefficient ", D^(1), "(x) [",1/s,"]")),
       pch = 20)
lines(mean_bin, mean_bin - mean_bin^3, col = 'red', lty = 2)
plotCI(mean_bin, D2, uiw = eD2, xlab = "x [a.u.]",
       ylab = expression(paste("Diffusion coefficient ", D^(2), "(x) [",1/s,"]")),
       pch = 20)
lines(mean_bin, mean_bin^2 + 1, col = 'red', lty = 2)
par(op)

# Fit polynomials to the estimated drift and diffusion coefficients
estD1 <- coef(lm(D1 ~ mean_bin + I(mean_bin^2) + I(mean_bin^3), weights = 1/eD1))
estD2 <- coef(lm(D2 ~ mean_bin + I(mean_bin^2), weights = 1/eD2))
detach(ests)

# Generate a time series from the reconstructed coefficients
set.seed(5048);
rec_x <- timeseries1D(N = 1e7, d10 = estD1[1], d11 = estD1[2], d12 = estD1[3],
                      d13 = estD1[4], d20 = estD2[1], d21 = estD2[2],
                      d22 = estD2[3], sf = sf)

# Compute and plot the increment PDFs for the original and the reconstructed
# time series for four different time lags
op <- par(no.readonly = TRUE)
par(mar = c(4.2, 5.4, 0.5, 0.5))
par(mfrow = c(1, 1));
plot(1, 1, log = "y", type = "n", xlim = c(-11, 12), ylim = c(1e-8, 5),
     xlab = expression(Delta[tau]/sigma[Delta[tau]]), ylab = "density")

tau <- c(1, 10, 100, 1000)
for(i in 1:4) {
    delta <- diff(Ux, lag = tau[i])
    rec_delta <- diff(rec_x, lag = tau[i])
    den <- density(delta)
    den$x <- den$x/sd(delta, na.rm = TRUE)
    rec_den <- density(rec_delta)
    rec_den$x <- rec_den$x/sd(rec_delta, na.rm = TRUE)
    lines(den, col = i)
    lines(rec_den, lty = 2, col = i)
}
par(op)


##########################################
# Two dimensional example
##########################################

# Construct the drift matrices
D1_1 <- matrix(0, nrow = 4,ncol = 4)
D1_1[1, 2]  =  1
D1_2 <- matrix(0, nrow = 4,ncol = 4)
D1_2[2, 1]  =  0.02
D1_2[4, 1]  =  -1
D1_2[1, 2]  =  0.03
D1_2[3, 2]  =  -1

# Construct the diffusion matrices
g_11 <- matrix(0, nrow = 3, ncol = 3)
g_12 <- matrix(0, nrow = 3, ncol = 3)
g_21 <- matrix(0, nrow = 3, ncol = 3)
g_22 <- matrix(0, nrow = 3, ncol = 3)

# Generate the time series without noise (only for plotting)
set.seed(4711)
Ux1 <- timeseries2D(N = 1e6, 0.145, 0.0002, D1_1, D1_2, g_11, g_12, g_21,
                    g_22, sf = sf)

# Add the noise
g_11[1, 1] <- 0.0025
g_22[1, 1] <- 0.0025

# Generate time series with added noise
set.seed(4711)
Ux <- timeseries2D(N = 1e8, 0.145, 0.0002, D1_1, D1_2, g_11, g_12, g_21,
                   g_22, sf = sf)

# Plot trajectories of the time series without and with noise
op <- par(no.readonly = TRUE)
par(mar = c(4.2, 5.4, 0.5, 0.5))
par(mfrow = c(1, 2))
plot(Ux1[1:1e6, 1], Ux1[1:1e6, 2], xlab = expression(paste(X[1]," [a.u.]")),
     ylab = expression(paste(X[2]," [a.u.]")), t = 'l')
plot(Ux[1:1e5, 1], Ux[1:1e5, 2], xlab = expression(paste(X[1]," [a.u.]")),
     ylab = expression(paste(X[2]," [a.u.]")), t = 'l')
par(op)

# Estimate drift and diffusion coefficients
ests <- Langevin2D(Ux, bins, steps)
attach(ests)
# Compute the mids of the bins for plotting
bin_mids <- matrix(NA, nrow = 40, ncol = 2)
bin_mids[,1] <- U[1:40, 1] + diff(U[,1])/2
bin_mids[,2] <- U[1:40, 2] + diff(U[,2])/2

# Plot drift and diffusion coefficients
op <- par(no.readonly = TRUE)
par(mar = c(4.2, 5.4, 0.5, 0.5))
layout(matrix(c(1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5), 2, 6, byrow = TRUE))

pmat <- perspx(x = bin_mids[,1], y = bin_mids[,2], z = D1[,,1], xlab = "", ylab = "",
               zlab = "", ticktype = "detailed", theta = 30, main = "")
attr(pmat, "ranges") <- list(x = range(bin_mids[,1]), y = range(bin_mids[,2]),
                             z = range(D1[,,1], na.rm = TRUE))
mtext3d("X", pmat, expression(X[1]))
mpos<-plotrix:::get_axispos3d("Y", pmat, at = 0, dist = .3)
mpos[1,] <- -mpos[1,]
ptext3d(mpos[1,], mpos[2,], mpos[3,], pmat = pmat, xpd = NA, texts = expression(X[2]))
mtext3d("Z", pmat, expression(D[x]^(1)), at = 0)

pmat <- perspx(x = bin_mids[,1], y = bin_mids[,2], z = D1[,,2], xlab = "", ylab = "",
               zlab = "", ticktype = "detailed", theta = 30, main = "")
attr(pmat, "ranges") <- list(x = range(bin_mids[,1]), y = range(bin_mids[,2]),
                             z = range(D1[,,2], na.rm = TRUE))
mtext3d("X", pmat, expression(X[1]))
mpos<-plotrix:::get_axispos3d("Y", pmat, at = 0, dist = .3)
mpos[1,] <- -mpos[1,]
ptext3d(mpos[1,], mpos[2,], mpos[3,], pmat = pmat, xpd = NA, texts = expression(X[2]));
mtext3d("Z", pmat, expression(D[y]^(1)), at = 0)

pmat <- perspx(x = bin_mids[,1], y = bin_mids[,2], z = D2[,,1], xlab = "", ylab = "",
               zlab = "", ticktype = "detailed", theta = 30, main = "")
attr(pmat, "ranges") <- list(x = range(bin_mids[,1]), y = range(bin_mids[,2]),
                             z = range(D2[,,1], na.rm = TRUE))
mtext3d("X", pmat, expression(X[1]))
mpos<-plotrix:::get_axispos3d("Y", pmat, at = 0, dist = .3)
mpos[1,] <- -mpos[1,]
ptext3d(mpos[1,], mpos[2,], mpos[3,], pmat = pmat, xpd = NA, texts = expression(X[2]));
mtext3d("Z", pmat, expression(D[xx]^(2)), at = 0.0025)

pmat <- perspx(x = bin_mids[,1], y = bin_mids[,2], z = D2[,,3], xlab = "", ylab = "",
               zlab = "", ticktype = "detailed", theta = 30, main = "")
attr(pmat, "ranges") <- list(x = range(bin_mids[,1]), y = range(bin_mids[,2]),
                             z = range(D2[,,3], na.rm = TRUE))
mtext3d("X", pmat, expression(X[1]))
mpos<-plotrix:::get_axispos3d("Y", pmat, at = 0, dist = .3)
mpos[1,] <- -mpos[1,]
ptext3d(mpos[1,], mpos[2,], mpos[3,], pmat = pmat, xpd = NA, texts = expression(X[2]));
mtext3d("Z", pmat, expression(D[yy]^(2)), at = 0.00225)

pmat <- perspx(x = bin_mids[,1], y = bin_mids[,2], z = D2[,,2], xlab = "", ylab = "",
               zlab = "", ticktype = "detailed", theta = 30, main = "")
attr(pmat, "ranges") <- list(x = range(bin_mids[,1]), y = range(bin_mids[,2]),
                             z = range(D2[,,2], na.rm = TRUE))
mtext3d("X", pmat, expression(X[1]))
mpos<-plotrix:::get_axispos3d("Y", pmat, at = 0, dist = .3)
mpos[1,] <- -mpos[1,]
ptext3d(mpos[1,], mpos[2,], mpos[3,], pmat = pmat, xpd = NA, texts = expression(X[2]))
mtext3d("Z", pmat, expression(D[xy]^(2)), at = -2.5e-4)
par(op)
detach(ests)