FaST-LMM light

Lightweight implementation of the FaST-LMM algorithm for estimating the parameters $\beta$, $\sigma^2$ and $\delta$ in a linear mixed model or Gaussian process of the form

$$y \sim N\bigl(X\beta, \sigma^2(K + \delta I) \bigr)$$

with $y\in\mathbb{R}^n$ an observed response vector, $X\in\mathbb{R}^{n\times d}$ an observed matrix of $d$ covariates, $K\in\mathbb{R}^{n\times n}$ a kernel matrix and $I$ the identity matrix.

"Lightweight" refers to two aspects of this implementation:

1. The package is agnostic to any application domain. Whereas the original software and an existing julia implementation contain a lot of code that is specific to working with genetic data, FaSTLMMlight contains only functions that work with the variables defined in the equation above, and nothing else. In other words, the package provides generic functions that other packages can call to solve the parameter estimation problem in specific application domains (for an example, see SpatialOmicsGPs).

2. The package introduces a new, lightweight formulation of the one-parameter cost function (after eliminating $\beta$ and $\sigma^2$). Whereas the original algorithm involves the solution of a linear system in each evaluation of the cost function (to express $\beta$ as a function of $\delta$), FaSTLMMlight uses restricted maximum likelihood properties to avoid this step, resulting in a cost function amenable to automatic differentiation. This in turns avoids the need for any grid-based optimization. See the documentation pages "SVD of the fixed effects" and "FaST-LMM full rank" for details.

Roadmap

FaSTLMMlight will implement the FaST-LMM algorithms for full-rank kernel matrices and low-rank kernel matrix factorizations. I am also investigating whether and when the Fourier-space method of Greengard et al. can lead to further speed-up in cases where the kernel matrix $K$ is coming from a Gaussian process with translationally invariant kernal.

Currently completed:

• LMMs/GPs with full-rank kernel matrix
• LMMs/GPs with low-rank kernel matrix factorization
• Fourier-space method for translationally invariant kernels