### Concave Weighted Least Squares ###
# 
# This file contains various programs connected
# to least squares estimation of a concave regression
# function.
# 
# 
### [M, IsKnots] = ConcaveWLS(X,Y,W) ###
# computes for given vectors X, Y in R^n such that
#    X[1] < X[2] < ... < X[n]
# and a weight vector W in (0,Inf)^n a vector M in R^n
# minimizing the weighted sum of squares
#    sum(W * (M - Y)^2 )
# under the constraint that
#    (M[i] - M[i-1])/(X[i] - X[i-1])
#    >=  (M[i+1] - M[i])/(X[i+1] - X[i])
#    for 2 <= i < n .
# 
# If one wants to estimatie a concave function m : R --> R
# minimizing
#    sum_{i=1}^n W[i]*(Y[i] - m(X[i]))^2
# for arbitrary observations
# (X[1],Y[1]), (X[2],Y[2]), ..., (X[n],Y[n]),
# one should preprocess the data vectors X, Y and (the
# optional) W with
#
### [x,y,w] <- PrepareData(X,Y,W) ###
# 
# This yields vectors x, y, w of length m, where
# x[1] < x[2] < ... < x[m] are the elements of the set
# {X[1], X[2], ..., X[n]}, whereas w[j] and y[j] are the
# sum of all W[i] (the number of all i) and the weighted
# average of all Y[i] such that X[i] = x[j].
# 
# 
### ConcaveWLS.demo1() and ConcaveWLS.demo2() ###
# illustrate the underlying active-set procedure.
# 
# 
### Examples ###
# 
# source("ConcaveWLS.R")
# 
# n <- 20
# X <- 3*(0:n)/n
# Y <- 1 - exp(-X) + 0.2*rnorm(n+1)
# Res <- ConcaveWLS.demo1(X,Y)
# 
# n <- 200
# X <- 3*(0:n)/n
# Y <- 1 - exp(-X) + 0.2*rnorm(n+1)
# Res <- ConcaveWLS.demo2(X,Y)
# 
# n <- 2000
# X <- 3*(0:n)/n
# Y <- 1 - exp(-X) + 0.2*rnorm(n+1)
# Res <- ConcaveWLS(X,Y,display=TRUE)
# 
# n <- 1000
# X <- round(3*runif(n),digits=1)
# Y <- 1 - exp(-X) + 0.2*rnorm(n)
# D <- PrepareData(X,Y)
# Res <- ConcaveWLS(X=D$x,Y=D$y,W=D$w,display=FALSE)
# plot(X,Y)
# lines(D$x,Res$M,lwd=2,col="blue")
# 
# Lutz Duembgen, Dezember 5, 2013


PrepareData <- function(X,Y,W=rep(1,length(X)))
{
	ord <- order(X)
	xx <- c(X[ord],Inf)
	yy <- Y[ord]
	ww <- W[ord]
	n <- length(X)
	tmp <- (xx[1:n] < xx[2:(n+1)])
	m <- sum(tmp)
	x <- xx[tmp]
	w <- c(0,cumsum(ww)[tmp])
	w <- w[2:(m+1)] - w[1:m]
	y <- c(0,cumsum(ww*yy)[tmp])
	y <- (y[2:(m+1)] - y[1:m])/w
	return(list(x=x,y=y,w=w))
}


ConcWLSRegr <- function(X,Y,W=rep(1,length(X)),
	display=FALSE)
{
	n <- length(Y)
	
	if (sum(X[2:n] <= X[1:(n-1)]) > 0)
	{
		print("Need a vector X such that X[1] < X[2] < ... !")
		print("Use PrepareData() to preprocess your data.")
		return(NULL)
	}
	
	IsKnot <- rep(0,n)
	IsKnot[1] <- 1
	IsKnot[n] <- 1
	M <- LocalEstimate(X,Y,W,IsKnot)
	H <- LocalDirDeriv(X,M-Y,W,IsKnot)
	
	# Precision parameter:
	prec <- 10^(-9) * mean(abs(Y-M)) / sd(X)
	
	f.old  <- Inf
	f.curr <- sum(W*(Y-M)^2)
	iter <- 0
	if (display)
	{
		print(paste("Iteration ",
			as.character(iter),
			": f(M) = ",
			as.character(f.curr),
			sep=""))
	}
	while (max(H) > prec && f.curr < f.old)
	{
		iter <- iter + 1
		# extend the set of knots:
		IsKnot[min((1:n)[H == max(H)])] <- 1
		# compute new candidate:
		M_new <- LocalEstimate(X,Y,W,IsKnot)
		# check new candidate's concavity:
		Conv     <- LocalConvexities(X,M)
		Conv_new <- LocalConvexities(X,M_new)
		while (max(Conv_new) > prec)
		{
			# modify M_new and reduce the set of knots:
			JJ <- (1:n)[Conv_new > prec]
			t <- min(- Conv[JJ] / (Conv_new[JJ] - Conv[JJ]))
			M <- (1 - t)*M + t*M_new
			Conv <- LocalConvexities(X,M)
			IsKnot <- LocalKnots(Conv,prec)
			M_new <- LocalEstimate(X,Y,W,IsKnot)
			Conv_new <- LocalConvexities(X,M_new)
		} # end while
		f.old <- f.curr
		f.curr <- sum(W*(Y - M_new)^2)
		if (f.curr < f.old)
		{
			M <- M_new
			Conv <- Conv_new
			IsKnot <- LocalKnots(Conv,prec)
			H <- LocalDirDeriv(X,M-Y,W,IsKnot)
			if (display)
			{
				print(paste("Iteration ",
					as.character(iter),
				": f(M) = ",
				as.character(f.curr),
				sep=""))
			}
		}
	} # end while
	if (display)
	{
		par(mai=c(0.5,0.5,0.1,0.1))
		plot(X,Y)
		points(X[IsKnot==1],M[IsKnot==1],pch=16,col="blue")
		lines(X,M,lwd=2,col="blue")
	}
	return(list(M=M,IsKnot=IsKnot))
}


ConcWLSRegr.demo <- function(X,Y,W=rep(1,length(X)))
{
	n <- length(Y)
	
	if (sum(X[2:n] <= X[1:(n-1)]) > 0)
	{
		print("Need a vector X such that X[1] < X[2] < ... !")
		return(NULL)
	}

	### For display:
	par(mai=c(0.5,0.5,0.5,0.1),cex=1,cex.axis=1)
	###
	
	IsKnot <- rep(0,n)
	IsKnot[1] <- 1
	IsKnot[n] <- 1
	M <- LocalEstimate(X,Y,W,IsKnot)
	H <- LocalDirDeriv(X,M-Y,W,IsKnot)
	
	# Precision parameter:
	prec <- 10^(-9) * mean(abs(Y-M)) / sd(X)
	
	f.old  <- Inf
	f.curr <- sum(W*(Y-M)^2)
	iter <- 0
	while (max(H) > prec && f.curr < f.old)
	{
		iter <- iter + 1
		
		###
		tmp <- readline("Next plot ... ")
		plot(X,Y,
			main=paste("Current proposal (locally optimal), f(M) = ",
				as.character(f.curr),sep=""))
		points(X[IsKnot==1],M[IsKnot==1],pch=16,col="blue")
		lines(X,M,lwd=2,col="blue")
		###
		
		# extend the set of knots:
		IsKnot[min((1:n)[H == max(H)])] <- 1
		# compute new candidate:
		M_new <- LocalEstimate(X,Y,W,IsKnot)
		# check new candidate's concavity:
		Conv     <- LocalConvexities(X,M)
		Conv_new <- LocalConvexities(X,M_new)
		
		###
		tmp <- readline("next plot ... ")
		points(X[IsKnot==1],
			M_new[IsKnot==1],pch=16,col="magenta")
		lines(X,M_new,lwd=1,col="magenta")
		if (max(Conv_new) > prec)
		{
			tmp <- (Conv_new > prec)
			points(X[tmp],M_new[tmp],pch=16,col="red")
		}
		###
		
		while (max(Conv_new) > prec)
		{
			# modify M_new and reduce the set of knots:
			JJ <- (1:n)[Conv_new > prec]
			t <- min(- Conv[JJ] / (Conv_new[JJ] - Conv[JJ]))
			M <- (1 - t)*M + t*M_new
			Conv <- LocalConvexities(X,M)
			IsKnot <- LocalKnots(Conv,prec)
			
			###
			tmp <- readline("Next plot ... ")
			plot(X,Y,
				main=paste("Current proposal (within cone), f(M) = ",
					as.character(sum(W*(Y - M)^2)),sep=""))
			points(X[IsKnot==1],
				M[IsKnot==1],pch=16,col="blue")
			lines(X,M,lwd=2,col="blue")
			###
			
			M_new <- LocalEstimate(X,Y,W,IsKnot)
			Conv_new <- LocalConvexities(X,M_new)
			
			###
			tmp <- readline("next plot ... ")
			points(X[IsKnot==1],
				M_new[IsKnot==1],pch=16,col="magenta")
			lines(X,M_new,lwd=1,col="magenta")
			if (max(Conv_new) > prec)
			{
				tmp <- (Conv_new > prec)
				points(X[tmp],M_new[tmp],pch=16,col="red")
			}
			###

		} # end while
		f.old <- f.curr
		f.curr <- sum(W*(Y - M_new)^2)
		if (f.curr < f.old)
		{
			M <- M_new
			Conv <- Conv_new
			IsKnot <- LocalKnots(Conv,prec)
			H <- LocalDirDeriv(X,M-Y,W,IsKnot)
		}
	} # end while

	###
	tmp <- readline("Final plot! ")
	plot(X,Y,
		main=paste("Final solution, f(M) = ",
			as.character(f.curr),sep=""))
	points(X[IsKnot==1],M[IsKnot==1],pch=16,col="blue")
	lines(X,M,lwd=2,col="blue")
	###

	return(list(M=M,IsKnot=IsKnot))
}


LocalEstimate <- function(X,Y,W,IsKnot)
{
	n <- length(Y)
	JJ <- which(IsKnot == 1)
	p <- length(JJ)
	DD <- matrix(rep(0,n*p),nrow=n,ncol=p)
	DD[1,1] <- 1
	for (i in 1:(p-1))
	{
		a <- JJ[i]
		b <- JJ[i+1]
		tmp <- (X[(a+1):b] - X[a]) / (X[b] - X[a]);
		DD[(a+1):b,i]   <- 1 - tmp
		DD[(a+1):b,i+1] <-   tmp
	}
	DD <- diag(sqrt(W)) %*% DD
	M <- (DD %*% qr.solve(DD, sqrt(W)*Y)) / sqrt(W)
	return(M)
}


LocalDirDeriv <- function(X,MRes,W,IsKnot)
{
	n <- length(X)
	H <- rep(0,n)
	JJ <- (1:n)[IsKnot == 1]
	tmpcs <- cumsum(W * MRes)
	for (i in 1:(length(JJ)-1))
	{
		a <- JJ[i]
		b <- JJ[i+1]
		if (a < b-1)
		{
			tmpdx <- X[(a+1):(b-1)] - X[a:(b-2)]
			H[(a+1):(b-1)] = cumsum(tmpdx * tmpcs[a:(b-2)])
		}
	}
	return(H)
}


LocalConvexities <- function(X,M)
{
	n <- length(X)
	dXw <- X[3:n] - X[1:(n-2)]
	l1 <- (X[2:(n-1)] - X[1:(n-2)]) / dXw
	l2 <- 1 - l1
	dXw <- dXw * l1 * l2
	# Compute Conv:
	Conv <- rep(0,n)
	Conv[2:(n-1)] = (l1*M[3:n] + l2*M[1:(n-2)] - M[2:(n-1)]) / dXw
	return(Conv)
}


LocalKnots <- function(conv,prec)
{
	n <- length(conv)
	IsKnot <- rep(0,n)
	IsKnot[1] <- 1
	IsKnot[n] <- 1
	IsKnot[conv < - prec] <- 1
	return(IsKnot)
}