Randomly Pivoted Cholesky

This contains code to reproduce the numerical experiments for Randomly pivoted Cholesky: Practical approximation of a kernel matrix with few entry evaluations by Yifan Chen, Ethan N. Epperly, Joel A. Tropp, and Robert J. Webber and the forthcoming paper [Embrace rejection: Kernel matrix approximation by accelerated randomly pivoted Cholesky] by Ethan N. Epperly, Joel A. Tropp, and Robert J. Webber.

Citing this Repository

If you use our code in your work, please cite the following BibTeX entries:

@article{chen2023randomly,
  title={Randomly pivoted Cholesky: {Practical} approximation of a kernel matrix with few entry evaluations},
  author={Yifan Chen and Ethan N. Epperly and Joel A. Tropp and Robert J. Webber},
  journal={arXiv preprint arXiv:2207.06503},
  url = {https://arxiv.org/abs/2207.06503},
  year={2023}
}

@article{epperly2024embrace,
  title={Embrace rejection: {Kernel} matrix approximation by accelerated randomly pivoted {Cholesky}},
  author={Ethan N. Epperly and Joel A. Tropp and Robert J. Webber},
  journal={Manuscript in preparation},
  year={2024}
}

Algorithm

Randomly pivoted Cholesky (RPCholesky) is a fast, randomized algorithm for computing a low-rank approximation to a positive semidefinite matrix $\boldsymbol{A}$ (i.e., a symmetric matrix with nonnegative eigenvalues). We anticipate the algorithm is most useful for kernel and Gaussian process methods, where the matrix $\boldsymbol{A}$ is defined only implicitly by a kernel function which must be evaluated at great expense to read each entry of $\boldsymbol{A}$. The result of the algorithm is a rank $k$ (column) Nyström approximation $\boldsymbol{\hat{A}} = \boldsymbol{F}\boldsymbol{F}^*$ to the matrix $\boldsymbol{A}$, computed in $\mathcal{O}(k^2N)$ operations and only $(k+1)N$ entry evaluations of the matrix $\boldsymbol{A}$. In our experience, RPCholesky consistently provides approximations of comparable accuracy to other methods with fewer entry evaluations. By default, this paper uses the faster accelerated RPCholesky method, developed in followup work. Accelerated RPCholesky method produces the exact same random output as ordinary RPCholesky, but is faster.

Using RPCholesky

RPCholesky is implemented in the rpcholesky method in rpcholesky.py, and can be called as

nystrom_approximation = rpcholesky(A, num_pivots)

The input matrix A should be an AbstractPSDMatrix object, defined in matrix.py. A psd matrix stored as a numpy array can be made usable by rpcholesky by wrapping it in a PSDMatrix object:

nystrom_approximation = rpcholesky(PSDMatrix(ordinary_numpy_array), num_pivots)

The output of rpcholesky is a PSDLowRank object (defined in lra.py). From this object, one can obtain $\boldsymbol{F}$ (defining the Nyström approximation $\boldsymbol{\hat{A}} = \boldsymbol{FF}^*$), the pivot set $\mathsf{S}$, and the rows defining the Nyström approximation $A(\mathsf{S},:)$:

nystrom_approximation = rpcholesky(A, num_pivots)
F      = nystrom_approximation.get_left_factor()    # Nystrom approximation is F @ F.T
pivots = nystrom_approximation.get_indices() 
rows   = nystrom_approximation.get_rows()           # rows = A[pivots, :]

Running the experiments for the original RPCholesky paper

The first step to reproducing the experiments from the manuscript is to run the script

./setup.sh

which sets up the file structure, loads RLS and DPP samplers, and downloads the QM9 dataset for the KRR example. The data from the figures in the paper can produced by running the following scripts in the experiments/ folder, each of which has instructions for its individual use at a comment at the top:

1. experiments/comparison.py: compares the approximation error for different Nyström methods. Used to produce the left displays in Figure 1.
2. experiments/chosen.py: outputs the pivots chosen by different Nyström methods. Used to produce the right displays in Figure 1.
3. experiments/entries.py: outputs the entry evaluations for different Nyström methods. Used to produce Figure 2.
4. experiments/qm9_krr.py: performs kernel ridge regression on the QM9 dataset. Used to produce Figure 3.
5. experiments/cluster_biomolecule.py: performs spectral clustering on the alanine dipeptide dataset. Used to produce Figure 4.
6. experiments/timing.py: compares the timing of different Nyström methods.

Once the relevant Python scripts have been run, the figures from the paper can be generated from the relevant MATLAB scripts in experiments/matlab_plotting/.

Figure 4 in the manuscript was completely changed in revision. Figure 4 from earlier versions of the manuscript can be generated using the scripts experiments/cluster_letters.py and experiments/cluster_letters_plot.py.

Running the experiments for the accelerated RPCholesky paper

The experiments for the accelerated RPCholesky paper are found in the folder block_experiments/.