PCA User Story

[jerowe]

I'm going to start working on the PCA user story.

I've been back and forth with this on a few calls and internal discussion, but the general consensus seems to be that exploratory statistics and filtering should be separated out from the PCA functionality itself. It's already separated out in the codebase, but is grouped together in the docs as they would be together in a real world analysis.

Resources:

Scikit Allel PCA Docs
Tour of Scikit Allel
Fast PCA Blog Post

PCA analysis would be after LD pruning https://github.com/pystatgen/sgkit/issues/31

[eric-czech]

Hey @jerowe by upstream did you mean after LD pruning? We might have the wrong "stream" direction in mind but I wanted to verify that LD pruning would come first. It's a minor point for the PCA story though if you were going to start with real data, it might be a good idea to use locate_unlinked or sample with big fixed step windows (1 Mbp would do it if you wanted to try human data like 1000 genomes).

[jerowe]

@eric-czech yes I had the wrong stream. Those always get me. ;-)

PCA is after LD pruning.

The PCA functions in scikit-allel assume you've already done all the filtering, so I'm going along with that. I just thought it was worth mentioning that in an analysis it comes after so I'm following along with the user story.

[jerowe]

@alimanfoo @eric-czech (and whoever else wants to provide an opinion ;-) ) Questions!

I was digging into the PCA today. At first I was thinking just drop in dask_ml with some fancy doc string examples for those that are used to the , but then as I was looking at it I was wondering is that what we want?

Do the genotypes ever get big enough that we wouldn't want to just calculate in memory?

Also, it was decided on the last call that everyone likes the interface over PCA that allele provides, so keep that, right?

[jeromekelleher]

I think the consensus was that the interface that scikit-allel uses is pretty good, so we should just copy that in the first instance (@eric-czech - correct me if I'm wrong here!), with a straightforward implementation using upstream tools. Getting that and some unit test coverage into the library would be a great step in the right direction I think. Is this something you'd like to work into a PR @jerowe?

[eric-czech]

Makes sense to me @jerowe and @jeromekelleher! Adding a few more details:

I think dask-ml PCA is the right place to start and that genotypes will very often be bigger than memory so supporting the tall-skinny and fully chunked cases is a priority. By "fully chunked" I mean da.linalg.svd_compressed is used rather than da.linalg.svd internally.

As far as I can tell there is no reason not to use the dask-ml PCA estimator directly, except for pca.py#L270. That would only be an efficiency issue though so we can probably find a way to work around that later. There is also a minor discrepancy with the degrees of freedom used at pca.py#L285 and pca.py#L290 vs scikit-allele#randomized_pca though I don't think that's important at all.

BTW the PattersonScaler scales by binomial variance rather than empirical variance and that's a common approach in genetics with some nicer properties (i.e. StandardScaler) if you can assume HWE for the variants. I thought I'd share that because it wasn't clear to me early on why it's important. Otherwise it's basically identical to a StandardScaler.

If it helps, I also made this notebook a while ago using Xarray, dask-ml and the "Patterson" scaling logic for a fully-blockwise PCA: https://nbviewer.jupyter.org/github/related-sciences/gwas-analysis/blob/master/notebooks/platform/xarray/pca.ipynb. That doesn't wrap the scaling or PCA in a sklearn estimator though, and I believe we all want that.

I'll propose an end-state for PCA and scaling utilities, and perhaps @alimanfoo or @jeromekelleher can chime in if they disagree:

1. We create an actual subclass of BaseEstimator for the scaler and introduce scikit-learn as a dependency (as in rolling your own estimator). This should include the TransformerMixin too.
2. We use PCA like this:

from sgkit import PattersonScaler
from sklearn.pipeline import Pipeline
from dask_ml.decomposition import PCA
est = Pipeline([
  ('scaler', PattersonScaler()), 
  ('pca', PCA(n_components=2, svd_solver='randomized', random_state=1))
])
pcs = est.fit_transform(ds.call_genotype.sum(dim='ploidy'))
# `est` would also contain loadings + explained variance
# `scaler` would contain the MAF and binomial variance values needed for out-of-sample projection

# Out-of-sample projection
pcs_oos = est.transform(ds2.call_genotype.sum(dim='ploidy'))

If it becomes necessary, we could roll our own PCA estimator but it will still work within a Pipeline like that as long as we stick to subclassing BaseEstimator.

I'm not sure if want to still call it a PattersonScaler but that seems reasonable to me. HWEScaler would be another option.

---

p.s. An example showing how to do out-of-sample PCA projections with the "Patterson" model, that Hail calls hwe_normalized, is here. As far as I know, it should be possible to separate the scaling and transforming without needing any shared learned parameters that would force us to mash them into a single estimator.

[eric-czech]

Oh and one more thought on that -- I think this will be a lot for new users to learn. Much like the discussion on Xarray vs domain-specific wrapper classes, this will provide a lot of flexibility while introducing even more stuff to understand like scikit-learn pipelines. I would advocate for starting this way anyhow and then wrapping these up as individual functions or in a separate higher-level API later.

[jeromekelleher]

I would advocate for starting this way anyhow and then wrapping these up as individual functions or in a separate higher-level API later.

💯

[eric-czech]

And another thought -- I think it would be good if we could demonstrate something like this https://gist.github.com/eric-czech/8b6e0331f7e512f89cf009839e9f84ca (linked from dask/dask#6227 (comment)) using tskit/msprime instead. The Balding-Nichols dask array simulation function at the top is the simplest way I know of to create a dataset with multiple populations, but it would be cool to see an analog with tskit/msprime since it would support so much more simulation functionality.

[jerowe]

@jeromekelleher sure I'll work on this if no one else is. I have 10 hours / week on this, so I'd prefer more isolated functionality. PCA is fairly standalone. ;-) If anyone else wants to grab it that's fine too.

Is there a consensus on use the interface in scikit-allel vs sklearn Pipelines? (Edited - scikit-allel vs sklearn pipelines)

[eric-czech]

scikit vs sklearn

Assuming that was scikit-allele vs sklearn, I think we should use sklearn instead but they should be identical really since @alimanfoo had already done everything you'd need to for a custom estimator in PattersonScaler. IIRC the "_" suffix on variables set after a call to fit is crucial to allow for est.copy() to work correctly in cross-validation wrappers (so only the hyperparameters are copied). Alistair had that already so all that should be left is to make it descend from BaseEstimator instead of object and think through how it works on da.Array/xr.DataArray. I'm not sure why it isn't already a BaseEstimator since scikit-learn is a requirement for scikit-allele (which I guess means we're already depending on scikit-learn).

[jeromekelleher]

@jeromekelleher sure I'll work on this if no one else is. I have 10 hours / week on this, so I'd prefer more isolated functionality. PCA is fairly standalone. ;-)

It's all yours @jerowe, have fun!

[jerowe]

It's all yours @jerowe, have fun!

Always!