wSVD.Rmd

---
title: "wSVD"
author: "stephens999"
date: "2018-06-28"
output: workflowr::wflow_html
---

# Introduction

Motivated by the desire to apply SVD and related methods
to non-gaussian data (eg single cell data), I want to suggest
investigating "weighted" versions of SVD that allows each observation
to have its own measurement-error variance (in addition to any common variance). 
We already have this kind of idea in flash and mash, but methods like
softImpute and svd are potentially more scalable, and so it would be nice
to implement fast general versions of these.

The working model "rank k" model is
$$X = UDV' + Z + E$$
where $X$, $Z$ and $E$ are all $n \times p$ matrices, and $D$ is a $k \times k$ diagonal matrix.
The elements of $E$ are iid 
$$E_{ij} \sim N(0,\sigma^2=1/\tau)$$ 
and 
$$Z_{ij} \sim N(0,s^2_{ij})$$ 
where $s_{ij}$ are known.

Note: in softImpute (alternating least squares; ALS version) 
they replace $UDV'$ by $AB'$, but the basic idea is the same.
Also in softImpute they introduce an L2 penalty, which is a nice feature to have,
and which I think may not complicate things much here. (to be checked!)

Given $Z$ we note that the mle for $U,D,V$ is given by the SVD of ($X-Z$).
Following the usual EM idea, each iteration we can replace $Z$ with its expectation $\bar{Z} = E(Z | U,D,V)$ where $U,D,V$ are the current values of these parameters.
Then the M step becomes running SVD on $X-\bar{Z}$.

Given $U,D,V$ define residuals $R= X-UDV$. Then from the model
$R_{ij} | Z \sim N(Z_{ij}, \sigma^2)$. Then from standard
Bayesian analysis of Gaussians we have:
$$Z_{ij} | R \sim N(\mu_1,1/\tau_1)$$
where 
$$\mu_1 = \tau/\tau_1 R_{ij}$$
$$\tau_1 = \tau + 1/s_{ij}^2$$.

In particular the conditional mean of $Z$ needed for EM is:
$$\bar{Z}_{ij}= \tau/\tau_1 R_{ij}$$.

Note that in the special case $s_{ij}=\Inf$, which is like $X_{ij}$ is "missing",
this gives $\bar{Z}_{ij} = R_{ij}$, and when we plug that in to get a "new" value of $R$ we get
$R_{ij} = X_{ij}-\bar{Z}_{ij} = (UDV)_{ij}$.
That is, each iteration 

If we look in the softImpute code this is exactly what they use to deal with missing data. For example, line 49 of `simpute.als.R` is  
`xfill[xnas] = (U %*% (Dsq * t(V)))[xnas]`.

# Idea

Basically my idea is that we should be able to modify the softImpute code
by replacing this line (and similar lines involving xfill) with something
based on the above derivation...One advantage of this is that softImpute
already deals with ridge penalty, and is well documented and fast...

Alternatively we could just implement it ourselves as below without the ridge penalty...

# Code

I started coding an EM algorithm that
imputes $Z$ each iteration.  I haven't tested it, so there may be bugs ... but the
objective seems to increase. This code may or may not be useful to build on.

```{r}
wSVD = function(x,s,k,niter=100,tau=NULL){
  if(is.null(tau)){ #for now just estimate tau by residual variance from fit of first k svd
    x.svd = svd(x,k,k)
    tau = 1/mean((x - x.svd$u %*% diag(x.svd$d[1:k]) %*% t(x.svd$v))^2) 
  }
  n = nrow(x)
  p = ncol(x)
  z = matrix(0,nrow=n,ncol=p)
  sigma2 = rep(0,niter)
  obj = rep(0,niter)
  for(i in 1:niter){
    x.svd = svd(x-z,k,k) # could maybe replace this with a faster method to get top k pcs?
    R = x - x.svd$u %*% diag(x.svd$d[1:k]) %*% t(x.svd$v)
    tau1 = tau + 1/s^2 
    z =  (tau/tau1)*R
    sigma2[i] = 1/tau
    obj[i] = sum(dnorm(R, 0, sqrt(s^2+(1/tau)), log=TRUE))
  }
  return(list(svd = x.svd,sigma2=sigma2,obj=obj))
}
```

This example just runs it on constant $s$, so it should match regular svd.
I ran it with two different values of $\tau$, but I don't think $\tau$ affects the mle here...
```{r}
set.seed(1)
n = 100
p = 1000
s = matrix(1,nrow=n,ncol=p)
x = matrix(rnorm(n*p,0,s),nrow=n,ncol=p)
x.wsvd = wSVD(x,s,3,10,1)
plot(x.wsvd$obj)
x.wsvd2 = wSVD(x,s,3,10,1e6)
plot(x.wsvd2$svd$u[,1],svd(x,3,3)$u[,1])
plot(x.wsvd$svd$u[,1],svd(x,3,3)$u[,1])
```

Now this version is non-constant variance. At least the objective is increasing...
```{r}
s = matrix(rgamma(n*p,1,1),nrow=n,ncol=p)
x = matrix(rnorm(n*p,0,s),nrow=n,ncol=p)
x.wsvd = wSVD(x,s,30,100)
plot(x.wsvd$obj)
```


# Applying to Poisson data

Here is outline of how we might apply this to Poisson data.

Let $Y_{ij} \sim Poi(\mu_{ij})$ be observations.
Let $m_{ij}$ be a "current fitted value" for $\mu_{ij}$, for example
it could be the current value of $Y_{ij} - Z_{ij}$ if $Z_{ij}$ where $Z_{ij}$
is the estimated "measurement error" in the above.
Or for initialization $m_{ij}$ could be the mean of $X_{i\cdot}$ if $j$ indexes
genes. 

Basically the idea is to do a Taylor series expansion of the Poisson likelihood about $m_{ij}$. This leads to
- set $X_{ij} = \log(m_{ij}) + (Y_{ij}-m_{ij})/m_{ij}$
- set $s^2_{ij} = 1/m_{ij}$ 
which we could apply wSVD to. And then repeat...

# Other issues

- I think there is something of a literature on wSVD, but not much implemented.
Would be good to look further. 

- We would want to estimate the residual precision $\tau$ too. I believe there is
a EM update for that based on computing expected squared residuals (which will involve second moment of $Z_{ij}$, which is available).

- It would be nice if we could estimate the ridge penalty in softImpute by
maximum likelihood/variational approximation (as in flash/PEER) rather than having to do cross-validation.