# Principal warps computation and visualization

# This file contains the basic functions:
# (1) Fcomp_deform:  to evaluate the thin-plate spline on a set of target points
# (2) Fdraw_deform:  draw deformation grids
# Following these function definitions is simple example with code to
# visualize mapping of individual partial warps.

#library(MASS)

Fcomp_deform <- function(X, beta, Xtarg, xscaling=NA) {
  # Function to evaluate thin-plate spline specified by coefficients of partial warps, beta
  # Apply transformation to Xtarg, which may be the same as coordinates X,
  # assumed to have been standardized outside this function.
  # If an argument "xscaling" is provided, back transform to original coord system.
  n <- nrow(X)
  #  X <- scale(X)/sqrt(2*(n-1))

  h <- as.matrix(dist(X,diag=TRUE,upper=TRUE))
  K <- h^2 * log(h)
  diag(K) <- 0
  one <- rep(1,n)
  Gamma <- cbind(K,one,X)
  Gamma <- rbind(Gamma,c(one,0,0,0))
  Gamma <- rbind(Gamma,cbind(t(X),matrix(0,2,3)))
  Ginv <- solve(Gamma)
  Ginv <- (Ginv + t(Ginv))/2  # make exactly symmetric prior to eigen

  B <- Ginv[1:n,1:n]
  Beig <- eigen(B)
  g <- Beig$vectors
  l <- Beig$values

  g <- g[,order(l)]
  l <- l[order(l)]

  parWarpsSum<- matrix(0,nrow=nrow(X),ncol=ncol(X))
  for (j in 4:n) {
    parWarpsSum <- parWarpsSum + g[,j] %o% beta[,j-2]
  }

  # The affine part
  c <- c(0,0)
  A <- matrix(c(1,beta[1,1],0,1+beta[2,1]),byrow=X,2,2)

  # Calculating Y
  htarg <- as.matrix(dist(rbind(X,Xtarg),diag=TRUE,upper=TRUE))
  htarg <- htarg[1:n,(n+1):(n+nrow(Xtarg))]
  sigma <- htarg^2 * log(htarg)
  sigma[htarg==0] <- 0

  Y <- t( matrix(c,2,nrow(Xtarg)) +
      A %*% t(Xtarg) + t(parWarpsSum) %*% sigma )

  if (length(xscaling)==3) { # Back-transformation
    Yr <- (Y%*%t(xscaling$rotation))
    Y <- Yr * xscaling$scale 
    Y <- scale(Y, center = -xscaling$mean, scale=FALSE)
  }
  return(Y)
}


Fdraw_deform <- function(X,beta,xscaling=NA,drawsq=T,
      cexpts=1,colpts="red") {

  xlim <- range(pretty(X[,1]))
  ylim <- range(pretty(X[,2]))
  xseq <- seq(xlim[1],xlim[2],.05)
  yseq <- seq(ylim[1],ylim[2],.05)
  nx <- length(xseq)
  ny <- length(yseq)

  xygrid <- as.matrix(expand.grid(xseq,yseq))
  Y <- Fcomp_deform(X,beta,X,xscaling)
  Ygrid <- Fcomp_deform(X,beta,xygrid,xscaling)
  nx1 <- nx-1
  ny1 <- ny-1

  if (length(xscaling)==3) {
    Xr <- (X%*%t(xscaling$rotation))
    X <- Xr * xscaling$scale 
    X <- scale(X, center = -xscaling$mean, scale=FALSE) 
    xygridr <- (xygrid%*%t(xscaling$rotation))
    xygrid <- xygrid * xscaling$scale 
    xygrid <- scale(xygrid, center = -xscaling$mean, scale=FALSE) 
  }
  
  if(drawsq) {
  eqscplot(range(xygrid[,1]),range(xygrid[,2]),type="n",xlab="",ylab="")
  points(X,pch=16,cex=cexpts/2,col=colpts)
  text(X[,1],X[,2],cex=1)
  for (i in 1:ny) {
    lines(xygrid[(1+(i-1)*nx):((1+(i-1)*nx)+nx1),])
  }
  for (j in 1:nx) {
    lines(xygrid[seq((1+(j-1)),nx*ny,nx),])
  }
  }

  eqscplot(range(Ygrid[,1]),range(Ygrid[,2]),type="n",xlab="",ylab="")
  points(Y,pch=16,cex=cexpts,col=colpts)
  for (i in 1:ny) {
    lines(Ygrid[(1+(i-1)*nx):((1+(i-1)*nx)+nx1),])
  }
  for (j in 1:nx) {
    lines(Ygrid[seq((1+(j-1)),nx*ny,nx),])
  }
}




#### Example computations

# The labeling of the plots below uses the eigenvalues and
# eigenvectors computed above in Fcomp_deform, but I ran
# that code separately as I did not return the values from
# the eigendecomposition.

## initial data
#  X - matrix of initial coordinates
#  n - number of spatial coordinates
#  B - bending energy matrix
#  beta - matrix of coefficients affine / partial warp coefficients

X <- matrix(c(-.2, 1, -1, 0, 0, -1, 1, 0, 0, 0, .9, -.9, 1.1, -1.1),
            byrow=T,ncol=2)
X0 <- X
# Standardize by centering, scaling to centroid size = 1,
# and rotating to the principal axes
n <- nrow(X)
Xs <- scale(X)/sqrt(2*(n-1))
Xsvd <- svd(Xs)
# Here I will make sure diag elements of rotation matrix are positive
# by scaling by scaling columns by +/- 1.  This is to keep the
# rotation from flipping the orientation.
Xsvd$v[,1] <- Xsvd$v[,1] * sign(Xsvd$v[1,1]) # 
Xsvd$v[,2] <- Xsvd$v[,2] * sign(Xsvd$v[2,2])
Xt <- Xs %*% Xsvd$v
X <- Xt

# Compute bending energy marix and eigenvectors
h <- as.matrix(dist(X,diag=TRUE,upper=TRUE))
K <- h^2 * log(h)
diag(K) <- 0
one <- rep(1,n)
Gamma <- cbind(K,one,X)
Gamma <- rbind(Gamma,c(one,0,0,0))
Gamma <- rbind(Gamma,cbind(t(X),matrix(0,2,3)))
Ginv <- solve(Gamma)
Ginv <- (Ginv + t(Ginv))/2  # make exactly symmetric prior to eigen

B <- Ginv[1:n,1:n]
Beig <- eigen(B)
g <- Beig$vectors
l <- Beig$values

g <- g[,order(l)]
l <- l[order(l)]

# Coefficient matrix
# the first column  consists of the 'a' and 'b' parameters of the affine matrix;
# the last n-3 columns will be coefficients for the partial warps
beta <- matrix(c(-.5, 0, 0, 2, 0,
                   1, 1, 0, 0, 0),byrow=T,ncol=5)
beta <- matrix(c(-.5, .25, 0, 3, 0,
                   1, .25, .5, 0, 5),byrow=T,ncol=5)

Y <- Fcomp_deform(Xt,beta,Xt)

par(mfrow=c(1,2),mgp=c(2,.75,0),mar=c(4,3,2,1))
Fdraw_deform(Xt,beta,xscaling=F)

### Plot principal warp eigenvectors, direction arbitrary
pdf("Principal warp eigenvectors.pdf")
xscaling <- NA
par(mfrow=c(1,1))
#xlim <- range(pretty(X[,1]))
#ylim <- range(pretty(X[,2]))
xlim <- c(-1.4,1.4)  #xlim <-c(-0.7,0.7)
ylim <- c(-1.4,1.4)  #ylim <- c(-0.4,0.4)
xseq <- seq(xlim[1],xlim[2],.1)
yseq <- seq(ylim[1],ylim[2],.1)
nx <- length(xseq)
ny <- length(yseq)

xygrid <- as.matrix(expand.grid(xseq,yseq))
#Y <- Fcomp_deform(X,beta,X,xscaling)
#Ygrid <- Fcomp_deform(X,beta,xygrid,xscaling)
nx1 <- nx-1
ny1 <- ny-1

eqscplot(range(xygrid[,1]),range(xygrid[,2]),type="n",xlab="",ylab="",
         xaxt="n",yaxt="n")
points(X0,pch=16,cex=.25)
text(X0[,1],X0[,2],cex=2)
for (i in 1:ny) {
  lines(xygrid[(1+(i-1)*nx):((1+(i-1)*nx)+nx1),])
}
for (j in 1:nx) {
  lines(xygrid[seq((1+(j-1)),nx*ny,nx),])
}

alen <- .4
arrows(X0[,1],X0[,2],X0[,1]+0*g[,4],X0[,2]+sqrt(alen)*g[,4],col="red",
       length=.1,lwd=2)
arrows(X0[,1],X0[,2],X0[,1]+2*alen*g[,5],X0[,2]+2*alen*g[,5],col="dark green",
       length=.1,lwd=2)
arrows(X0[,1],X0[,2],X0[,1]+sqrt(alen)*g[,6],X0[,2]+0*g[,6],col="blue",
       length=.1,lwd=2)
arrows(X0[,1],X0[,2],X0[,1]+alen*g[,7],X0[,2]-alen*g[,7],col="purple",
       length=.1,lwd=2)
legend("topright",col=c("red","dark green","blue","purple"),
                        lwd=c(2,2,2,2),
                        legend=c(paste("PW 1, eval",round(l[4],2)),
                                 paste("PW 2, eval",round(l[5],2)),
                                 paste("PW 3, eval",round(l[6],2)),
                                 paste("PW 4, eval",round(l[7],2))))
title(main="Principal Warps: Bending energy matrix eigenvectors")
title(sub=paste(
  paste("Evect 1: ",paste(round(g[,4],2),collapse=", ")),"\n",
  paste("Evect 2: ",paste(round(g[,5],2),collapse=", ")),"\n",
  paste("Evect 3: ",paste(round(g[,6],2),collapse=", ")),"\n",
  paste("Evect 4: ",paste(round(g[,7],2),collapse=", ")),"\n"))
  
dev.off()

##  Play around, pushing each warp toward folding by increasing
##  magnitude of deformation via beta coefficients
pdf("Partial Warps illustrations.pdf")
par(mfrow=c(2,2),mgp=c(2,.75,0))

beta <- matrix(c(0,0,0,0,0,
                 .75,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="Affine warp 'y'")
beta <- matrix(c(.75,0,0,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="Affine warp 'x'")

beta <- matrix(c(0,0,0,0,0,
                 1.5,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="Affine warp 'y', scale 2")
beta <- matrix(c(1.5,0,0,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="Affine warp 'x', scale 2")

# PW 1
beta <- matrix(c(0,0,0,0,0,
                 0,1,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 1 'y'")
beta <- matrix(c(0,1,0,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 1 'x'")

# PW 1 scale 2
beta <- matrix(c(0,0,0,0,0,
                 0,2,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 1 'y', scale 2")
beta <- matrix(c(0,1,0,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 1 'x', scale 2")

# PW 2
beta <- matrix(c(0,0,0,0,0,
                 0,0,1.25,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 2 'y'")
beta <- matrix(c(0,0,1.25,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 2 'x'")

# PW 2 scale 2
beta <- matrix(c(0,0,0,0,0,
                 0,0,2.5,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 2 'y', scale 2")
beta <- matrix(c(0,0,2.5,0,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 2 'x', scale 2")


# PW 3
beta <- matrix(c(0,0,0,0,0,
                 0,0,0,2,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 3 'y'")
beta <- matrix(c(0,0,0,2,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 3 'x'")

# PW 3 scale 2
beta <- matrix(c(0,0,0,0,0,
                 0,0,0,4,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 3 'y', scale 2")
beta <- matrix(c(0,0,0,4,0,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 3 'x', scale 2")

# PW 4
beta <- matrix(c(0,0,0,0,0,
                 0,0,0,0,3),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 4 'y'")
beta <- matrix(c(0,0,0,0,3,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 4 'x'")

# PW 4 scale 2
beta <- matrix(c(0,0,0,0,0,
                 0,0,0,0,6),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 4 'y', scale 2")
beta <- matrix(c(0,0,0,0,6,
                 0,0,0,0,0),byrow=T,ncol=5)
Fdraw_deform(Xs,beta,drawsq=F)
title(main="PW 4 'x', scale 2")

dev.off()   # End, partial warps illustrations